--- license: cc-by-nc-nd-4.0 ---
branchsbm

Branched Schrödinger Bridge Matching

Sophia Tang1·Yinuo Zhang2·Alexander Tong3·Pranam Chatterjee4

1 University of Pennsylvania   2 Duke-NUS Medical School   3 Mila, Quebec AI Institute  

4 Duke University

arXiv
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.
branchsbm

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.
branchsbm

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.
branchsbm

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.
branchsbm

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} } ```