Multiscale Supervised Unbalanced Optimal Transport Flow Matching

TL;DR

MUST-FM is a scalable, multiscale framework for unbalanced optimal transport that incorporates hierarchical data and transition priors to improve trajectory inference in large single-cell datasets.

Contribution

It introduces MUST-FM, a novel multiscale supervised UOT method that leverages hierarchical structures and priors for efficient, accurate single-cell trajectory modeling.

Findings

Reduces computational cost significantly.

Achieves robust, biologically meaningful trajectories.

Enables dynamic modeling at atlas scale.

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.