Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Branched Schrödinger Bridge Matching

Sophia Tang, Yinuo Zhang, Alexander Tong, Pranam Chatterjee

TL;DR

BranchSBM parameterizes multiple time-dependent velocity fields and growth processes, enabling the representation of population-level divergence into multiple terminal distributions and is essential for tasks involving multi-path surface navigation, modeling cell fate bifurcations from homogeneous progenitor states, and simulating diverging cellular responses to perturbations.

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.

Branched Schrödinger Bridge Matching

TL;DR

BranchSBM parameterizes multiple time-dependent velocity fields and growth processes, enabling the representation of population-level divergence into multiple terminal distributions and is essential for tasks involving multi-path surface navigation, modeling cell fate bifurcations from homogeneous progenitor states, and simulating diverging cellular responses to perturbations.

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.

Paper Structure

This paper contains 92 sections, 16 theorems, 94 equations, 9 figures, 11 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Schrödinger Bridge Problem
  4. Generalized Schrödinger Bridge Problem
  5. Branched Schrödinger Bridge Matching
  6. Unbalanced Conditional Stochastic Optimal Control
  7. Unbalanced Generalized Schrödinger Bridge Problem
  8. Unbalanced Conditional Stochastic Optimal Control (CondSOC)
  9. BranchSBM: Sum of Unbalanced CondSOC Problems
  10. Branched Generalized Schrödinger Bridge Problem
  11. Branched Conditional Stochastic Optimal Control
  12. Learning BranchSBM Using Neural Networks
  13. Branched Neural Interpolant Optimization
  14. Learning the Energy-Minimizing Branching Dynamics
  15. Branched Energy Loss
  16. ...and 77 more sections

Key Result

Proposition 3.1

Suppose the marginal density can be decomposed as $p_t(X_t)=\int p_t(X_t|\boldsymbol{x}_0, \boldsymbol{x}_{1})\pi_{0, 1}(d\boldsymbol{x}_0, d\boldsymbol{x}_{1})$, where $\pi_{0, 1}$ is a fixed joint coupling of the data. Then, we can identify the optimal drift $u^\star_t$ and growth $g^\star_t$ that where $w_{t}= w_0+\int_0^tg_{s}(X_{s})ds$ is the time-dependent weight initialized at $w^\star_0$,

Figures (9)

  • Figure 1: Branched Schrödinger Bridge Matching(A) Stage 1 trains a correction term that learns the optimal interpolant conditioned on endpoints (B) Stage 2 and 3 trains a separate flow and growth network for each branch independently (C) Stage 4 jointly optimizes the flow and growth networks to minimize the energy, mass, and matching loss.
  • Figure 2: Plot of weight (left) and energy (right) calculated with (\ref{['loss:Energy']}) of each branch over time $t\in [0,1]$. Mass is transferred from the primary branch to branch 1, and both converge to the target weight of $0.5$ at $t=1$. Both plots represent the average over trajectories from samples in the validation set.
  • Figure 4: Application of BranchSBM on Modeling Differentiating Single-Cell Population Dynamics. Mouse hematopoiesis scRNA-seq data is provided for three time points $t_0,t_1, t_2$. (A) Simulated states (top) and trajectories (bottom) at time $t_1$ using single-branch SBM. (B) Simulated states with BranchSBM at $t_1$ ($t=0.5$) and (C)$t_2$ ($t=1$). (D) Learned trajectories over the interval $t\in [t_0,t_2]$ on validation samples.
  • Figure 5: Results for Trametinib Perturbation Modeling with BranchSBM.(A) Gene expression data of DMSO control ($t=0$) and cells after treatment with $5\mu M$ Trametinib ($t=0$) with three distinct endpoints (purple, turquoise, and pink). (B) The simulated endpoints of the top 50 PCs at $t=1$ on the validation data for each branch. (C) The evolution of cumulative energy across $t\in [0,1]$ calculated as (\ref{['loss:Energy']}) along each branched trajectory after Stage 3 (growth with fixed drift) and Stage 4 (joint) training. (D) The evolution of mass across $t\in [0,1]$ along each branched trajectory with target weights of $w_{1, 0}=0.603$, $w_{1, 1}=0.255$ and $w_{1, 2}=0.142$.
  • Figure 6: Results for Clonidine Perturbation Modeling.(A) Gene expression data of DMSO control (set to $t=0$) and cell states (set to $t=1$) after Clonidine perturbation with two distinct endpoints (pink and purple). (B) The simulated trajectories for single-branch SBM on the top 50 PCs with both clusters. All samples take the low-energy path without reaching the second cluster. (C) The simulated endpoints of the top 50, 100, and 150 PCs at $t=1$ on the validation data for each branch.
  • ...and 4 more figures

Theorems & Definitions (25)

  • Proposition 3.1: Unbalanced Conditional Stochastic Optimal Control
  • Proposition 3.2: Branched Conditional Stochastic Optimal Control
  • Remark 3.1
  • Proposition 4.1: Solving the GSB Problem with Stage 1 and 2 Training
  • Proposition 4.2: Existence of Optimal Growth Functions
  • Definition A.1: Reciprocal Class
  • Definition A.2: Markovian Projection
  • Definition A.3: Schrödinger Bridge
  • Proposition C.1: Unbalanced Conditional Stochastic Optimal Control
  • Proposition C.1: Branched Conditional Stochastic Optimal Control
  • ...and 15 more