Branched Schrödinger Bridge Matching

Sophia Tang¹ · Yinuo Zhang² · Alexander Tong³ · Pranam Chatterjee^{4^{¹ University of Pennsylvania
² Duke-NUS Medical School
³ Mila, Quebec AI Institute

⁴ Duke University}}

Predicting the intermediate trajectories between an initial and target distribution is a central problem in generative modeling. Existing approaches, such as flow matching and Schrödinger Bridge Matching, effectively learn mappings between two distributions by modeling a single stochastic path. However, these methods are inherently limited to unimodal transitions and cannot capture branched or divergent evolution from a common origin to multiple distinct outcomes. To address this, we introduce Branched Schrödinger Bridge Matching (BranchSBM), a novel framework that learns branched Schrödinger bridges. BranchSBM parameterizes multiple time-dependent velocity fields and growth processes, enabling the representation of population-level divergence into multiple terminal distributions. We show that BranchSBM is not only more expressive but also essential for tasks involving multi-path surface navigation, modeling cell fate bifurcations from homogeneous progenitor states, and simulating diverging cellular responses to perturbations.

Experiments

1. Branched LiDAR Surface Navigation

First, we evaluate BranchSBM for navigating branched paths along the surface of a 3-dimensional LiDAR manifold, from an initial distribution to two distinct target distributions.

Figure 3: Application of BranchSBM on Learning Branched Paths on a LiDAR Manifold.

2. Modeling Differentiating Single-Cell Population Dynamics

BranchSBM is uniquely positioned to model single-cell population dynamics where a homogeneous cell population (e.g., progenitor cells) differentiates into several distinct subpopulation branches, each of which independently undergoes growth dynamics. We demonstrate this capability on mouse hematopoiesis data.

Figure 4: Application of BranchSBM on Modeling Differentiating Single-Cell Population Dynamics.

3. Modeling Drug-Induced Perturbation Responses

Predicting the effects of perturbation on cell state dynamics is a crucial problem for therapeutic design. In this experiment, we leverage BranchSBM to model the trajectories of a single cell line from a single homogeneous state to multiple heterogeneous states after a drug-induced perturbation. We demonstrate that BranchSBM is capable of capturing the dynamics of high-dimensional gene expression data and learning branched trajectories that accurately reconstruct diverging perturbed cell populations.

First, we modeled two branches to two divergent subpopulations in the Clonidine-perturbed cells from the initial control DMSO-treated cells with BranchSBM and compared with single-branch SBM.

Figure 5: Results for Clonidine Perturbation Modeling with BranchSBM.

Finally, we used BranchSBM to model three branched trajectories in the Trametinib-perturbed cells from the initial control DMSO-treated cells.

Figure 6: Results for Trametinib Perturbation Modeling with BranchSBM.

Citation

If you find this repository helpful for your publications, please consider citing our paper:

@article{tang2025branchsbm,
  title={Branched Schrödinger Bridge Matching},
  author={Tang, Sophia and Zhang, Yinuo and Tong, Alexander and Chatterjee, Pranam},
  journal={arXiv preprint arXiv:2506.09007},
  year={2025}
}