Branched Schrödinger Bridge Matching

Sophia Tang, Yinuo Zhang, Alexander Tong, Pranam Chatterjee

Keywords: Schrödinger bridges, branched generative modeling, stochastic optimal control, unbalanced optimal transport, flow matching, trajectory inference, stochastic processes, probabilistic transport, multimodal distributions, dynamical systems

TL;DR: BranchSBM extends Schrödinger Bridge Matching to branched and unbalanced trajectories, enabling neural modeling of multimodal dynamical systems such as cell fate bifurcations and perturbation responses.

Abstract: 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 modes. 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.

Primary Area: applications to physical sciences (physics, chemistry, biology, etc.)

Submission Number: 5258