Multiscale Supervised Unbalanced Optimal Transport Flow Matching
Pith reviewed 2026-05-20 19:53 UTC · model grok-4.3
The pith
MUST-FM scales unbalanced optimal transport to atlas-scale single-cell datasets by using hierarchical structure and optional lineage priors.
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
Core claim
MUST-FM is a simulation-free framework that scales UOT by leveraging hierarchical data structure in single-cell experiments and supports an optional supervised formulation that incorporates transition priors such as cell lineages to guide the learning of displacement fields and mass variations, thereby reducing computational overhead while producing robust and biologically meaningful trajectory inference on atlas-scale datasets.
What carries the argument
Multiscale supervised unbalanced optimal transport flow matching, which processes data at multiple hierarchical levels to approximate transport plans and mass variations without simulation and optionally conditions the flows on known transition priors.
Load-bearing premise
Single-cell datasets contain reliable hierarchical annotations and accurate transition priors that can be used to guide the model without introducing bias or discarding important information.
What would settle it
Apply MUST-FM and a standard UOT baseline to the same large single-cell atlas with recorded cell lineages; if the new method shows no meaningful reduction in runtime or fails to recover the known trajectories at comparable accuracy, the scalability claim does not hold.
Figures
read the original 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.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The manuscript introduces Multiscale Supervised Unbalanced Optimal Transport Flow Matching (MUST-FM), a simulation-free framework that scales unbalanced optimal transport (UOT) for single-cell trajectory inference by leveraging hierarchical data annotations and an optional supervised formulation that incorporates transition priors such as cell lineages to learn displacement fields and mass variations. The central claim is that this multiscale approach reduces computational overhead relative to standard UOT approximations while producing robust, biologically meaningful results on atlas-scale datasets.
Significance. If the claims hold, the work would be significant for computational single-cell biology by making principled UOT-based modeling of birth-death dynamics feasible at scale, where existing approximations do not exploit hierarchical structure or priors. The simulation-free and annotation-driven design is a practical strength that could enable new applications in dynamic atlas modeling.
major comments (2)
- [Abstract and §3] Abstract and §3 (Method): The central scalability claim rests on the assumption that available hierarchical annotations decompose the UOT problem into scales that preserve the true displacement fields and mass variations. No quantitative sensitivity analysis or worst-case bound on approximation error is provided for the case where the hierarchy reflects annotation artifacts rather than dynamical structure, which directly risks the robustness claim under hierarchy mismatch.
- [§5] §5 (Experiments): The statement that MUST-FM 'reduces computational overhead while achieving robust results' is load-bearing for the contribution, yet the abstract and summary provide no specific baselines, runtime/memory metrics, dataset scales, or accuracy measures (e.g., trajectory reconstruction error or Wasserstein distance) against standard UOT solvers or flow-matching variants. This prevents verification of the claimed improvement.
minor comments (2)
- [§3] Clarify the precise form of the supervised loss that incorporates transition priors and how it is combined with the multiscale UOT objective.
- [§2] Add a short discussion of related multiscale OT or hierarchical flow-matching methods to better situate the novelty.
Simulated Author's Rebuttal
We thank the referee for the constructive and detailed feedback. We address each major comment below and describe the revisions we will make to improve clarity and strengthen the claims.
read point-by-point responses
-
Referee: [Abstract and §3] The central scalability claim rests on the assumption that available hierarchical annotations decompose the UOT problem into scales that preserve the true displacement fields and mass variations. No quantitative sensitivity analysis or worst-case bound on approximation error is provided for the case where the hierarchy reflects annotation artifacts rather than dynamical structure, which directly risks the robustness claim under hierarchy mismatch.
Authors: We agree that robustness under imperfect hierarchies is an important consideration. While deriving a general worst-case theoretical bound would require substantial additional analysis beyond the current scope, we will add an empirical sensitivity study in the revised manuscript. This will involve systematically perturbing hierarchical labels on the experimental datasets and reporting the resulting changes in learned displacement fields, mass variation estimates, and downstream trajectory metrics. revision: partial
-
Referee: [§5] The statement that MUST-FM 'reduces computational overhead while achieving robust results' is load-bearing for the contribution, yet the abstract and summary provide no specific baselines, runtime/memory metrics, dataset scales, or accuracy measures (e.g., trajectory reconstruction error or Wasserstein distance) against standard UOT solvers or flow-matching variants. This prevents verification of the claimed improvement.
Authors: We will revise the abstract to explicitly state key quantitative results from §5, including dataset scales (e.g., number of cells), runtime and memory reductions relative to standard UOT solvers, and accuracy metrics such as trajectory reconstruction error and Wasserstein distances against flow-matching baselines. We will also add a concise summary table of these metrics to make the improvements immediately verifiable. revision: yes
- Deriving a rigorous worst-case theoretical bound on approximation error for arbitrary mismatches between the provided hierarchy and the underlying dynamical structure.
Circularity Check
No significant circularity; framework presented as independent construction
full rationale
The paper introduces MUST-FM as a novel simulation-free framework that leverages hierarchical annotations and optional transition priors to scale UOT. No equations or steps in the abstract or described method reduce by construction to fitted outputs or self-citations; the multiscale solver and supervised loss are algorithmic contributions that take hierarchy as external input rather than deriving it from the model's own predictions. The derivation chain remains self-contained against external benchmarks such as standard UOT solvers, with no load-bearing self-citation chains or ansatz smuggling identified.
Axiom & Free-Parameter Ledger
axioms (1)
- domain assumption Single-cell data contain hierarchical annotations and known transition priors that can be leveraged without loss of critical information.
Lean theorems connected to this paper
-
IndisputableMonolith/Cost/FunctionalEquation.leanwashburn_uniqueness_aczel unclear?
unclearRelation between the paper passage and the cited Recognition theorem.
Multiscale Supervised Unbalanced Optimal Transport Flow Matching (MUST-FM) ... hierarchical mask M^(l) ... OET objective (11) ... traveling Gaussian path ... CUFM loss (14)
-
IndisputableMonolith/Foundation/RealityFromDistinction.leanreality_from_one_distinction unclear?
unclearRelation between the paper passage and the cited Recognition theorem.
WFR metric ... Dirac-to-Dirac geodesic (2)–(4)
What do these tags mean?
- matches
- The paper's claim is directly supported by a theorem in the formal canon.
- supports
- The theorem supports part of the paper's argument, but the paper may add assumptions or extra steps.
- extends
- The paper goes beyond the formal theorem; the theorem is a base layer rather than the whole result.
- uses
- The paper appears to rely on the theorem as machinery.
- contradicts
- The paper's claim conflicts with a theorem or certificate in the canon.
- unclear
- Pith found a possible connection, but the passage is too broad, indirect, or ambiguous to say the theorem truly supports the claim.
Reference graph
Works this paper leans on
-
[1]
Ao Chen, Sha Liao, Mengnan Cheng, Kailong Ma, Liang Wu, Yiwei Lai, Xiaojie Qiu, Jin Yang, Jiangshan Xu, Shijie Hao, et al. Spatiotemporal transcriptomic atlas of mouse organogenesis using dna nanoball-patterned arrays.Cell, 185(10):1777–1792, 2022
work page 2022
-
[2]
Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming.Cell, 176(4):928–943, 2019
work page 2019
-
[3]
Grace XY Zheng, Jessica M Terry, Phillip Belgrader, Paul Ryvkin, Zachary W Bent, Ryan Wilson, Solongo B Ziraldo, Tobias D Wheeler, Geoff P McDermott, Junjie Zhu, et al. Massively parallel digital transcriptional profiling of single cells.Nature communications, 8(1):14049, 2017
work page 2017
-
[4]
Optimal transport for single-cell and spatial omics.Nature Reviews Methods Primers, 4(1):58, 2024
Charlotte Bunne, Geoffrey Schiebinger, Andreas Krause, Aviv Regev, and Marco Cuturi. Optimal transport for single-cell and spatial omics.Nature Reviews Methods Primers, 4(1):58, 2024
work page 2024
-
[5]
Zhenyi Zhang, Yuhao Sun, Qiangwei Peng, Tiejun Li, and Peijie Zhou. Integrating dynamical systems modeling with spatiotemporal scrna-seq data analysis.Entropy, 27(5), 2025
work page 2025
-
[6]
Deciphering cell-fate trajectories using spatiotemporal single-cell transcriptomic data
Zhenyi Zhang, Zihan Wang, Yuhao Sun, Jiantao Shen, Qiangwei Peng, Tiejun Li, and Peijie Zhou. Deciphering cell-fate trajectories using spatiotemporal single-cell transcriptomic data. npj Systems Biology and Applications, 2025
work page 2025
-
[7]
Mapping cells through time and space with moscot.Nature, pages 1–11, 2025
Dominik Klein, Giovanni Palla, Marius Lange, Michal Klein, Zoe Piran, Manuel Gander, Laetitia Meng-Papaxanthos, Michael Sterr, Lama Saber, Changying Jing, et al. Mapping cells through time and space with moscot.Nature, pages 1–11, 2025
work page 2025
-
[8]
Trajecto- rynet: A dynamic optimal transport network for modeling cellular dynamics
Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. Trajecto- rynet: A dynamic optimal transport network for modeling cellular dynamics. InInternational conference on machine learning, pages 9526–9536. PMLR, 2020
work page 2020
-
[9]
Guillaume Huguet, Daniel Sumner Magruder, Alexander Tong, Oluwadamilola Fasina, Manik Kuchroo, Guy Wolf, and Smita Krishnaswamy. Manifold interpolating optimal-transport flows for trajectory inference.Advances in neural information processing systems, 35:29705–29718, 2022
work page 2022
-
[10]
Jiaqi Zhang, Erica Larschan, Jeremy Bigness, and Ritambhara Singh. scNODE: gen- erative model for temporal single cell transcriptomic data prediction.Bioinformatics, 40(Supplement_2):ii146–ii154, 09 2024. 10
work page 2024
-
[11]
Neural ordinary differential equations.Advances in neural information processing systems, 31, 2018
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations.Advances in neural information processing systems, 31, 2018
work page 2018
-
[12]
On robustness of neural ordinary differential equations.arXiv preprint arXiv:1910.05513, 2019
Hanshu Yan, Jiawei Du, Vincent YF Tan, and Jiashi Feng. On robustness of neural ordinary differential equations.arXiv preprint arXiv:1910.05513, 2019
-
[13]
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. InThe Eleventh International Conference on Learning Representations, 2023
work page 2023
-
[14]
Building Normalizing Flows with Stochastic Interpolants
Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants.arXiv preprint arXiv:2209.15571, 2022
work page internal anchor Pith review Pith/arXiv arXiv 2022
-
[15]
Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based genera- tive models with minibatch optimal transport.Transactions on Machine Learning Research,
-
[16]
Expert Certification
-
[17]
Dominik Klein, Théo Uscidda, Fabian Theis, and Marco Cuturi. Genot: Entropic (gromov) wasserstein flow matching with applications to single-cell genomics.Advances in Neural Information Processing Systems, 37:103897–103944, 2024
work page 2024
-
[18]
Yutong Sha, Yuchi Qiu, Peijie Zhou, and Qing Nie. Reconstructing growth and dynamic trajectories from single-cell transcriptomics data.Nature Machine Intelligence, 6(1):25–39, 2024
work page 2024
-
[19]
Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal transport in competition with reaction: The hellinger–kantorovich distance and geodesic curves.SIAM Journal on Mathematical Analysis, 48(4):2869–2911, 2016
work page 2016
-
[20]
Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard. An interpolat- ing distance between optimal transport and fisher–rao metrics.Foundations of Computational Mathematics, 18(1):1–44, 2018
work page 2018
-
[21]
A new optimal transport distance on the space of finite Radon measures
Stanislav Kondratyev, Léonard Monsaingeon, and Dmitry V orotnikov. A new optimal transport distance on the space of finite radon measures.arXiv preprint arXiv:1505.07746, 2016
work page internal anchor Pith review Pith/arXiv arXiv 2016
-
[22]
stvcr: spatiotemporal dynamics of single cells
Qiangwei Peng, Peijie Zhou, and Tiejun Li. stvcr: spatiotemporal dynamics of single cells. Nature Methods, pages 1–12, 2026
work page 2026
-
[23]
Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport
Zhenyi Zhang, Tiejun Li, and Peijie Zhou. Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport. InThe Thirteenth International Conference on Learning Representations, 2025
work page 2025
-
[24]
Variational regularized un- balanced optimal transport: Single network, least action
Yuhao Sun, Zhenyi Zhang, Zihan Wang, Tiejun Li, and Peijie Zhou. Variational regularized un- balanced optimal transport: Single network, least action. InThe Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025
work page 2025
-
[25]
Unbalancedness in neural monge maps improves unpaired domain translation
Luca Eyring, Dominik Klein, Théo Uscidda, Giovanni Palla, Niki Kilbertus, Zeynep Akata, and Fabian J Theis. Unbalancedness in neural monge maps improves unpaired domain translation. InThe Twelfth International Conference on Learning Representations, 2024
work page 2024
-
[26]
arXiv preprint arXiv:2505.13413(2025)
Dongyi Wang, Yuanwei Jiang, Zhenyi Zhang, Xiang Gu, Peijie Zhou, and Jian Sun. Joint velocity-growth flow matching for single-cell dynamics modeling.arXiv preprint arXiv:2505.13413, 2025
-
[27]
WFR-FM: Simulation-free dynamic unbalanced optimal transport
Qiangwei Peng, Zihan Wang, Junda Ying, Yuhao Sun, Qing Nie, Lei Zhang, Tiejun Li, and Peijie Zhou. WFR-FM: Simulation-free dynamic unbalanced optimal transport. InThe Fourteenth International Conference on Learning Representations, 2026
work page 2026
-
[28]
Xinyu Wang, Ruoyu Wang, Qiangwei Peng, Peijie Zhou, and Tiejun Li. Wfr-mfm: One-step inference for dynamic unbalanced optimal transport.arXiv preprint arXiv:2601.20606, 2026
-
[29]
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport.Advances in neural information processing systems, 26, 2013. 11
work page 2013
-
[30]
Multiscale strategies for computing optimal transport
Samuel Gerber and Mauro Maggioni. Multiscale strategies for computing optimal transport. Journal of Machine Learning Research, 18(72):1–32, 2017
work page 2017
-
[31]
A multiscale approach to optimal transport
Quentin Mérigot. A multiscale approach to optimal transport. InComputer graphics forum, volume 30, pages 1583–1592. Wiley Online Library, 2011
work page 2011
-
[32]
Meyer Scetbon and Marco Cuturi. Low-rank optimal transport: Approximation, statistics and debiasing.Advances in Neural Information Processing Systems, 35:6802–6814, 2022
work page 2022
-
[33]
Meyer Scetbon, Michal Klein, Giovanni Palla, and Marco Cuturi. Unbalanced low-rank optimal transport solvers.Advances in Neural Information Processing Systems, 36:52312–52325, 2023
work page 2023
-
[34]
Max Sommerfeld, Jörn Schrieber, Yoav Zemel, and Axel Munk. Optimal transport: Fast proba- bilistic approximation with exact solvers.Journal of Machine Learning Research, 20(105):1–23, 2019
work page 2019
-
[35]
The single-cell transcriptional landscape of mammalian organogenesis.Nature, 566(7745):496–502, 2019
Junyue Cao, Malte Spielmann, Xiaojie Qiu, Xingfan Huang, Daniel M Ibrahim, Andrew J Hill, Fan Zhang, Stefan Mundlos, Lena Christiansen, Frank J Steemers, et al. The single-cell transcriptional landscape of mammalian organogenesis.Nature, 566(7745):496–502, 2019
work page 2019
-
[36]
Jeffrey A Farrell, Yiqun Wang, Samantha J Riesenfeld, Karthik Shekhar, Aviv Regev, and Alexander F Schier. Single-cell reconstruction of developmental trajectories during zebrafish embryogenesis.Science, 360(6392):eaar3131, 2018
work page 2018
-
[37]
A lineage-resolved molec- ular atlas of c
Jonathan S Packer, Qin Zhu, Chau Huynh, Priya Sivaramakrishnan, Elicia Preston, Hannah Dueck, Derek Stefanik, Kai Tan, Cole Trapnell, Junhyong Kim, et al. A lineage-resolved molec- ular atlas of c. elegans embryogenesis at single-cell resolution.Science, 365(6459):eaax1971, 2019
work page 2019
-
[38]
A single-cell transcriptomic atlas characterizes ageing tissues in the mouse.Nature, 583(7817):590–595, 2020
work page 2020
-
[39]
Molecular architecture of the developing mouse brain.Nature, 596(7870):92–96, 2021
Gioele La Manno, Kimberly Siletti, Alessandro Furlan, Daniel Gyllborg, Elin Vinsland, Alejan- dro Mossi Albiach, Christoffer Mattsson Langseth, Irina Khven, Alex R Lederer, Lisa M Dratva, et al. Molecular architecture of the developing mouse brain.Nature, 596(7870):92–96, 2021
work page 2021
-
[40]
Supervised optimal transport.SIAM Journal on Applied Mathematics, 82(5):1851–1877, 2022
Zixuan Cang, Qing Nie, and Yanxiang Zhao. Supervised optimal transport.SIAM Journal on Applied Mathematics, 82(5):1851–1877, 2022
work page 2022
-
[41]
On the translocation of masses
Leonid V Kantorovich. On the translocation of masses. InDokl. Akad. Nauk. USSR (NS), volume 37, pages 199–201, 1942
work page 1942
-
[42]
Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the monge-kantorovich mass transfer problem.Numerische Mathematik, 84(3):375–393, 2000
work page 2000
-
[43]
Jean-David Benamou. Numerical resolution of an “unbalanced” mass transport problem.ESAIM: Mathematical Modelling and Numerical Analysis, 37(5):851–868, 2003
work page 2003
-
[44]
Alessio Figalli. The optimal partial transport problem.Archive for rational mechanics and analysis, 195(2):533–560, 2010
work page 2010
-
[45]
Luis A Caffarelli and Robert J McCann. Free boundaries in optimal transport and monge-ampere obstacle problems.Annals of mathematics, pages 673–730, 2010
work page 2010
-
[46]
Dest-ot: Alignment of spatiotemporal transcriptomics data.Cell Systems, 16(2), 2025
Peter Halmos, Xinhao Liu, Julian Gold, Feng Chen, Li Ding, and Benjamin J Raphael. Dest-ot: Alignment of spatiotemporal transcriptomics data.Cell Systems, 16(2), 2025
work page 2025
-
[47]
Modeling complex system dynamics with flow matching across time and conditions
Martin Rohbeck, Edward De Brouwer, Charlotte Bunne, Jan-Christian Huetter, Anne Biton, Kelvin Y Chen, Aviv Regev, and Romain Lopez. Modeling complex system dynamics with flow matching across time and conditions. InThe Thirteenth International Conference on Learning Representations, 2025
work page 2025
-
[48]
Action matching: Learning stochastic dynamics from samples
Kirill Neklyudov, Rob Brekelmans, Daniel Severo, and Alireza Makhzani. Action matching: Learning stochastic dynamics from samples. InInternational conference on machine learning, pages 25858–25889. PMLR, 2023. 12
work page 2023
-
[49]
Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard. Unbalanced optimal transport: Dynamic and kantorovich formulations.Journal of Functional Analysis, 274(11):3090–3123, 2018
work page 2018
-
[50]
Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal entropy-transport problems and a new hellinger–kantorovich distance between positive measures.Inventiones mathematicae, 211(3):969–1117, 2018
work page 2018
-
[51]
Mingyue Wang, Qinan Hu, Zhencheng Tu, Lingshi Kong, Tengxiang Yu, Zihan Jia, Yuetian Wang, Jiajun Yao, Rong Xiang, Zhan Chen, et al. A drosophila single-cell 3d spatiotem- poral multi-omics atlas unveils panoramic key regulators of cell-type differentiation.Cell, 188(17):4734–4753, 2025
work page 2025
-
[52]
Spatiotemporal modeling of molecular holograms
Xiaojie Qiu, Daniel Y Zhu, Yifan Lu, Jiajun Yao, Zehua Jing, Kyung Hoi Min, Mengnan Cheng, Hailin Pan, Lulu Zuo, Samuel King, et al. Spatiotemporal modeling of molecular holograms. Cell, 187(26):7351–7373, 2024
work page 2024
-
[53]
Colin Kern, Qingquan Zhang, Yifan Lu, Jacqueline Eschbach, Zehua Zeng, Elie N Farah, Chu- Yi Tai, Kaifu Yang, Ignatius Jenie, Fenyong Yao, et al. Merfish+, a large-scale, multi-omics spatial technology resolves the molecular holograms of the 3d human developing heart.bioRxiv, pages 2025–11, 2025
work page 2025
-
[54]
Pot: Python optimal transport.Journal of Machine Learning Research, 22(78):1–8, 2021
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, et al. Pot: Python optimal transport.Journal of Machine Learning Research, 22(78):1–8, 2021
work page 2021
-
[55]
Caleb Weinreb, Alejo Rodriguez-Fraticelli, Fernando D Camargo, and Allon M Klein. Lin- eage tracing on transcriptional landscapes links state to fate during differentiation.Science, 367(6479):eaaw3381, 2020
work page 2020
-
[56]
Kevin R Moon, David Van Dijk, Zheng Wang, Scott Gigante, Daniel B Burkhardt, William S Chen, Kristina Yim, Antonia van den Elzen, Matthew J Hirn, Ronald R Coifman, et al. Vi- sualizing structure and transitions in high-dimensional biological data.Nature biotechnology, 37(12):1482–1492, 2019. 13 A Appendix A A.1 Implicit Finest-Level Semi-Coupling Sampling...
work page 2019
discussion (0)
Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.