File size: 4,835 Bytes
5565829 985ef16 5565829 553d0ea 5565829 bbf8516 5565829 8223ba6 5565829 985ef16 5565829 985ef16 5565829 8c235ae 5565829 985ef16 5565829 985ef16 5565829 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
---
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> <b>·</b> 
<a href="https://www.linkedin.com/in/yinuozhang98/" target="_blank">Yinuo Zhang</a><sup>2</sup> <b>·</b> 
<a href="https://www.alextong.net/" target="_blank">Alexander Tong</a><sup>3</sup> <b>·</b> 
<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  
<sup>2</sup> Duke-NUS Medical School  
<sup>3</sup> Mila, Quebec AI Institute  
</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}
}
``` |