Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots
About
Inferring cellular trajectories from destructive snapshots is complicated by the challenges of stochasticity and non-conservative mass dynamics such as cell proliferation and apoptosis. Existing unbalanced Optimal Transport (OT) methods treat mass as a continuous fluid, performing inference at the population level. However, this macroscopic view often fails to capture the discrete, jump-like nature of birth-death events at single-cell resolution, which is essential for understanding lineage branching and fate decisions. We present Unbalanced Schr\"odinger Bridge (USB), a simulation-free framework for learning underlying dynamics that effectively integrates both stochastic and unbalanced effects which also models the discrete, jump-like birth-death dynamics at single-cell resolution. Theoretically, USB provides a tractable solution to the Branching Schr\"odinger Bridge (BSB) problem, offering a rigorous microscopic interpretation where individual cells undergo both Brownian motion and discrete birth-death jumps. Technically, the method implements an efficient solver by introducing a simulation-free training objective that effectively scales to high-dimensional omics data. Empirically, we demonstrate on both simulated and real-world datasets that USB not only achieves trajectory reconstruction performance better than or comparable to deterministic baselines but also uniquely enables realistic discrete simulation of birth-death dynamics at single-cell resolution.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Reconstruction of discrete branching dynamics | 5D EB t=1 | W10.331 | 24 | |
| Reconstruction of discrete branching dynamics | CITE 50D (t=3) | W19.421 | 24 | |
| Trajectory Interpolation | EB 50D (held-out time points) | Mean W16.988 | 24 | |
| Trajectory reconstruction | Gaussian Mixtures 1000D | W1 Distance2.136 | 18 | |
| Discrete Branching Dynamics Reconstruction | Dyngen t=1 | W1 Distance0.109 | 15 | |
| Discrete Branching Dynamics Reconstruction | Dyngen (t=2) | W1 Distance0.093 | 15 | |
| Discrete Branching Dynamics Reconstruction | Dyngen t=3 | W1 Distance0.18 | 15 | |
| Discrete Branching Dynamics Reconstruction | Dyngen t=4 | W1 Distance0.142 | 15 | |
| Discrete Branching Dynamics Reconstruction | EMT 10D reduced (t=1) | W1 Score0.1831 | 12 | |
| Discrete Branching Dynamics Reconstruction | EMT 10D reduced (t=2) | W1 Distance0.2159 | 12 |