README.md

---
license: cc-by-nc-nd-4.0
---
<div align="center">
  <img src="branchsbm.png" alt="branchsbm" width="1000" height="300">
</div>

<h1 align='center'>Branched Schrödinger Bridge Matching</h1>

<div align="center">
  <a href="https://sophtang.github.io/" target="_blank">Sophia Tang</a><sup>1</sup>&ensp;<b>&middot;</b>&ensp;
  <a href="https://www.linkedin.com/in/yinuozhang98/" target="_blank">Yinuo Zhang</a><sup>2</sup>&ensp;<b>&middot;</b>&ensp;
  <a href="https://www.alextong.net/" target="_blank">Alexander Tong</a><sup>3</sup>&ensp;<b>&middot;</b>&ensp;
  <a href="https://www.chatterjeelab.com/" target="_blank">Pranam Chatterjee</a><sup>4<sup>
  <br>
  <p style="font-size: 16px;">
  <sup>1</sup> University of Pennsylvania &emsp; 
  <sup>2</sup> Duke-NUS Medical School &emsp; 
  <sup>3</sup> Mila, Quebec AI Institute &emsp; 
  </p>
  <p style="font-size: 16px;">
  <sup>4</sup> Duke University
  </p>
</div>
    
<div align="center">
 <a href="https://arxiv.org/abs/2506.09007"><img src="https://img.shields.io/badge/Arxiv-2506.09007-red?style=for-the-badge&logo=Arxiv" alt="arXiv"/></a>

</div>

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.

<div align="center">
  <img src="lidar.png" alt="branchsbm" width="900" height="300">
</div>
<p align="center"><em>Figure 3: Application of BranchSBM on Learning Branched Paths on a LiDAR Manifold.</em></p>

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

<div align="center">
  <img src="mouse.png" alt="branchsbm" width="900" height="300">
</div>
<p align="center"><em>Figure 4: Application of BranchSBM on Modeling Differentiating Single-Cell Population Dynamics.</em></p>

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

<div align="center">
  <img src="clonidine.png" alt="branchsbm" width="900" height="300">
</div>
<p align="center"><em>Figure 5: Results for Clonidine Perturbation Modeling with BranchSBM.</em></p>

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

<div align="center">
  <img src="trametinib.png" alt="branchsbm" width="900" height="300">
</div>
<p align="center"><em>Figure 6: Results for Trametinib Perturbation Modeling with BranchSBM.</em></p>

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