Branched Schr\"odinger Bridge Matching

TL;DR

Branched Schr"odinger Bridge Matching (BranchSBM) introduces a new framework for modeling complex, multi-path stochastic processes, enabling the capture of diverging trajectories in generative modeling and biological applications.

Contribution

The paper proposes BranchSBM, a novel method that learns branched Schr"odinger bridges with multiple velocity fields, extending existing models to handle multi-modal, divergent distributions.

Findings

Enables modeling of branched, multi-path trajectories.

Improves representation of biological cell fate bifurcations.

Applicable to complex generative and biological systems.

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\"odinger 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\"odinger Bridge Matching (BranchSBM), a novel framework that learns branched Schr\"odinger 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.