Branched Schrödinger Bridge Matching
Authors: Sophia Tang
Yinuo Zhang
Alexander Tong
Pranam Chatterjee
Venue: ICLR 2026 Poster (2026)
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
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
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.

Published as a conference paper at ICLR 2026
BRANCHED SCHRÖDINGER BRIDGE MATCHING
Sophia Tang', Yinuo Zhang", Alexander Tong", Pranam Chatterjee 14.1
1 "Department of Computer and Information Science, University of Pennsylvania "Center Computational Biology. Duke-NUS Medical School, 3 "AITHYRA 'Department of Bioengineering, University of Pennsylvania
Corresponding author: pranam @ seas.upenn.cc.
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 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 (BranchSB.M), a novel frame- work that learns branched Schrödinger bridges. BranchSBM parameterizes multiple dimedeeenneet 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.
INTRODUCTION
Tasks like crowd navigation and modeling cell-state transitions under perturbation involve learning transport map between two empirically observed endpoint distributions, rather than a noisy prior required for denoising diffusion (Austin et al., 2021) and flow matching (Lipman et al., 2023). The Schrödinger Bridge (SB) (Schrödinger, 1931) problem seeks to identify an optimal stochastic map between a pair of endpoint distributions that minimizes the Kullback-LLeibler (KL.) divergence to an underlying reference process. Schrödinger Bridge Matching (SBM) solves the SB problem by parameterizing drift field that matches a mixture of conditional stochastic bridges between endpoint pairs that each minimize the KL. divergence from a known reference process. Typically, SBM assumes conservation of mass from the initial to the target distribution, which fails to capture dynamical population behaviors such as growth and destruction of mass, commonly seen in single- cell population data. Furthermore, prior methods focus transporting samples from an initial distribution to the target distributions via independent stochastic trajectories, without accounting for branching dynamics (Liu et al., 2023a; Tong et al., 2024b; Theodoropoulos et al., 2024; Liu et al., 2022; De Bortoli et al., 2021), which are susceptible to mode collapse when the system follows branched trajectory that diverges toward multiple distinct target modes.
The notion of branching is central to many real-world systems. For example, when a homogeneous cell population undergoes a perturbation such as gene knockouts or drug treatments, it frequently induces fate bifurcation as the cell population splits into multiple phenotypically distinct outcomes or commits to divergent cell fates (Shalem et al., 2014; Zhang et al., 2025a). These trajectories are observable in single-cell RNA sequencing (scRNA-seq) data, where each subpopulation independently evolves and undergoes growth or contraction along its trajectory toward a distinct terminal state (Dixit al., 2016 6). In this work, we introduce Branched Schrödinger Bridge Matching (BranchSBM), framework for solving the branched Schrödinger bridge problem by parameterizing diverging velocity fields and branch-specific growth rates, which together define set of conditional stochastic bridges from the common source to multiple terminal modes.
Our main contributions be summarized in the following points:

---

Published as a conference paper at ICLR 2026
1. We define the Branched Generalized Schrödinger Bridge problem and introduce BranchSBM. novel matching framework that learns optimal branched trajectories from an initial distribution To to multiple target distributions {
2. We derive the Branched Conditional Stochastic Optimal Control (CondSOC) problem the sum of Unbalanced CondSOC objectives and leverage a stale training algorithm to learn the optimal branching drift and growth fields that transport mass along a branched trajectory.
3. We demonstrate the unique capability of BranchSBM to model dynamic branching trajectories while matching multiple target distributions across various problems, including 3D navigation over LiDAR manifolds (Section 5.1), modeling differentiating single-cell population dynamics (Section 5.2), and predicting heterogeneous cell states after perturbation (Section 5.3).
2 PRELIMINARIES
Schrödinger Bridge Problem Given a reference probability path measure Q. the Schrödinger Bridge (SB) problem aims to find an optimal path measure PSB that minimizes the Kullback-Leiblee (KI.) divergence with Q while satisfying the boundary distributions Po To and P1 m.
pss min [L Po To. P1 mT}
where Q is commonly defined as standard Brownian motion. For an extended background and formal definition of Schrödinger Bridges, refer Definition A.3 and Appendix A.
Generalized Schrödinger Bridge Problem The solution to the standard SB problem minimizes the kinetic energy of the conditional drift term un (X,) that preserves the endpoints drawn from the coupling (vo. 21) ~ To,1 defined as
min / Ep. (X)] s.t.
ddx, = (X + adB; [o To, X1 ~T
(2)
where dB, is standard d-dimensional Brownian motion. To define more complex systems where the optimal dynamics cannot be accurately captured by minimizing the standard squared-Fuclidean in entropic OT (Vargas et al., 2021), the Generalized Schrödinger Bridge (GSB) problem introduces an additional non-linear state-cost (x,) (Chen et al., 2021b; Chen & Georgiou, 2016; Liu et al., 2022). The minimization objective becomes
Euclidean cost
min /' the
ue (x.) + v.(x,) di s.t.
sdx, = (X,)) + adB, (xo ~ To, X1 ~ m
Epp
where the state cost can also be interpreted as the potential energy of the system at state X
3 BRANCHED SCHRÖDINGER BRIDGE MATCHING
A key challenge in trajectory matching is reconstructing multi-modal marginals, particularly when modes diverge along distinct dynamical paths. Existing Schrödinger bridge and flow matching frame- works approximate multi-modal distributions by simulating many independent particle trajectories, which are susceptible to mode collapse, with particles concentrating on dominant high-density modes or traversing only low-energy intermediate paths. To address this challenge, we introduce Branched Schrödinger Bridge Matching (BranchSBM), framework that learns set of diverging velocity fields to reconstruct multi-modal target distributions while simultaneously learning growth networks that allocate mass across branches. Guided by time-dependent potential V₁, BranchSBM captures diverging, eneryy-minimzzing dynamics without requiring intermediate- dineitttttieeeee supervision and can generate the full branched evolution from single initial sample.
3.1 UNRALANCED CONDITIONAL. STOCHASTIC OPTIMAL. CONTROL.
Unbalanced Generalized Schrödinger Bridge Problem Extending the definition of the General- ized Schrödinger Bridge (GSB) problem in Equation define the Unbalanced GSB problem by
2

---

Published as a conference paper at ICLR 2026
B
TAR
he (0) El sلتلجي
Lng(t)) | v((xa)))
Gy Comeallly Crennn (9)
Figure I: Branched Schrödinger Bridge Matching (A) Stage 1 trains correction term that Icarns the optimal interpolant conditioned endpoints (B) Stage 2 and trains separate flow and growth network for each branch independently (C) Stage jointly optimizes flow and growth networks minimize the energy, mass, and matching loss.
scaling the minimization objective by time-dependent weight wi wo + So 9s (X,)ds that evolves according to a time-varying growth rate (X.)) I2d x [0, 1] + R.
dX: (x(x))dtt + odB; No X1 wa. wi
min / Ep.
u4 (X.)11 va(xa) 2
(x.))d +v((x.
s.t.
(4)
Unbalanced Conditional Stochastic Optimal Control (CondSOC) Now, show that solve the Unbalanced GSB problem as an Unbalanced CondSOC problem where the optimal drift u, and growth ge minimize the expectation of the objective in (4) conditioned pairs of endpoints.
Proposition 3.1 (Unbalanced Conditional Stochastic Optimal Control). Suppose the marginal æ1 To. (dxo, dx11), where T0,1 is fixed
density can be decomposed as (X,) joint coupling the data. Then, can identify the optimal drift ot and growth 9 that solves the Unbalanced GSB problem (4) by minimizing the Unbalanced Conditional Stochastic Optimal Control objective given by
((TRRie [("lluaXa m)) 1 vi(X.) wr (x.) dt
min Bi )moTo,1 me:ge
(5)
S.f. dX,- = (e x adB,, No= ma, X1 = #1 wo wō, 10 = wi
(6)
where ue wo f/ gx sd is the time-dependent weight initialized at wh u is the drift. gr the growth rate, and To,1 is the weighted coupling of paired endpoints (æo, wy æi, x1, no.
The proof is provided in Appendix C.I. This defines the objective for tractably solve the Unbalanced GSB problem by conditioning on a finite set of endpoint pairs in the dataset.
3.2 BRANCHSBM: SUM OF UNBALANCED CONDSOC PROBLEMS
Branched Generalized Schrödinger Bridge Problem define the Branched GSB problem as minimizing the sum of Unbalanced GSB problems across all branches. All mass begins along primary path indexed it = 0 with initial weight 1. Over t [0, 1). mass is transferred across И secondary branches with initial weight and target weight wi,a such that it minimizes the objective defined
Given the Unbalanced GSB problem, we
min (1. M dX (x..)d adB, No no X1,6 R1,k; wo.k 8:=o, Wi.k wi.
J EEpL [1
( vi(Xoo) weo + Σ Ep.
1 (Xe,,)11 Vo(Xe.))
I
wa. Jat
k=l
(7)

---

Published as a conference paper at ICLR 2026
When total mass across branches is conserved, enforce k= Wi, = 1 for all t [0, 1]. which constrains the growth rates such that (t,o(Xx.o) + Σ (x,,(Xx, = This ensures that mass lost from the primary branch (when 9r,0 < 0) is redistributed among the secondary branches (where 9t.k > for some E (1...., ......,K))) K)). The primary branch evolves from an initial weight of according to wr, 1 + So go ds and the K secondary branches grow from the primary branch from weight O according to Wi,k = Jo (X,,)))
Branched Conditional Stochastic Optimal Control Following a similar procedure as shown for the Unbalanced GSB problem, we can reformulate the Branched GSB problem solving the Branched CondSOC problem where we optimize a set of parameterized drift {ut, / and growth (gt, 1K == networks by minimizing the energy of the conditional trajectories between paired samples (moo {2) ( ~ (
Proposition 3.2 (Branched Conditional Stochastic Optimal Control). For each branch, let (x,,( =
p (X To, Tu (dæu, dæ1.k), where T0,1, is the joint coupling distri-
bution of samples 26 ~ To from the initial distribution and 2i,k ~ 11.k from the kth target distribution. Then, we can identify the set of optimal drift and growth functions {ut & F k k=0 that solve the Branched GSB problem (7) by minimizing the sum of Unbalanced CondSOC objectives given by
L. Epraaoo ( [x..,.xx...)1+++++ Vo(Xe.o) wro Jai
min (1 KE
E E(
I ( (X + () [: {Epronn k=1
K
Vi(Xe We. dt
(8)
S.l. dKt.k = (X.x(()))dtt + odBi, Xo 2o. X1,x = 21,40 Wo,k = δe-o, W1.k wi,
(9)
where Wt,0 1 So (x,0(XX,,)) dss the weight of the primary paths initialized at and W = Is.k sdss are the weights of the K secondary branches initialized at O.
The proof is given in Appendix C.2. This defines the objective for to tractably solve the Branched GSB problem in Section 4 by conditioning on a discrete set of branched endpoint pairs in the dataset.
Remark 3.1. When 9t.0 = and St.k for all (xr t) R [0, 1) and l E (1,..., K}, then the Branched CondSOC problem is the solution to the single path GSB problem.
LEARNING BRANCHSBM USING NEURAL NETWORKS
Given an initial data distribution mo and K- target distributions ( BranchSBM aims to learn the optimal neural drift and growth fields that approximate the solution of the Branched CondSOC problem in Proposition 3.2. Our framework leverages multi-stage training algorithm that enables stable and scalable learning of branched dynamics from data snapshots.
4.1 BRANCHED NEURAL INTERPOLANT OPTIMIZATION
Since the optimal trajectory under the state cost Vy(X,) Rd x [0, + follows a non-linear cost manifold, given pair of endpoints (æg, 21,4) T0,1,k, train neural path interpolant PP.. (I., 11.k) Rd x Rd x [0. 1] - Rd that defines the intermediate state 2t,,, and velocity Lt.nkk = ,, at time t. which minimizes (3.2). We note that the neural interpolant has been introduced previously for single-pate interpolants in Kapuániak al. (2024); Neklyudov et al. (2023).
We define & to be bounded at the endpoints as given by
Lt.nk = t)æu + tæ + t(1 - t) (wo. 11,k)
(10)
k. 11.6 xo t(1 b) Fo, x1,e) (1 2t) 21,k)
(11)
4

---

Published as a conference paper at ICLR 2026
To optimize Pt.,y (Zo. 2i,k) such that predicts the eneryy-minimin trajectory, minimize the trajectory loss Luaj defined as
Cusj (1)
Σ (a mi -o.... k=0
2 u m k (.( dt
(12)
After convergence, Stage 1 returns the network Pi.,, (Zo, 11,a)) that generates the optimal conditional velocity , which defines the matching objective in Stage In Stage 2, we parameterize a set of neural drift fields uL,, (2t,,) IRd x [0, 1] IR that generates the mixture of bridges defined in Stage by minimizing the conditional flow matching loss (Lipman al., 2023; Tong et al., 2024a).
Lnww (0)
( | - (
(13)
Proposition 4.1 (Solving the GSB Problem with Stage I and 2 Training). Stage I and Stage 2 training yield the optimal drift ut (Xx) that generates the optimal marginal probability distribu- tion pt (X,) that solves the GSB problem in (3).
The proof is provided in Appendix C.3. Since the drift for each branch uf, (Xt,k) are trained independently in Stage 2, we can extend this result across all K branches and conclude that the sequential Stage 1 and Stage 2 training procedures yields the optimal set of drifts { J=o that generate the optimal probability paths { )/K that solves the GSB problem for each branch.
1K
4.2 LEARNING THE ENERGY- MINIMIIIGG BRANCHING DYNAMICS
solve the Branched CondSOC problem defined in Proposition 3.2, we minimize a branched energy loss Lenergy defined as
Branched Energy Loss
Leneryy (0, 6) -/- BipLatte +1 VV((x.))) w7o
Mee (Xt,)) vi(Xeu)
wit
pristary uajectory
16 branches
S.L. w.o = 1 S/'' G. (o,((,,,) W f (Xa,())))
(14)
At time the primary path has weight 1, and the K branches have weights 0. Over t [0, 1]. the weight of the primary path changes according (X,,)) and supplies mass the branches, which grow at rates Si (X4,,) (Lemma C.2). Intuitively, the branched energy loss optimizes the branching growth rates such that they are non-zero when branching is favored over the primary path.
Weight Matching Loss We define weight matching loss Lmarch that aims to minimize the dif- ference between the predicted weights of each branch 1, obtained by integrating the growth function F.A, (Xt,k) over E [0, 1). and the true weights of each terminal distribution {
"EEE BpI. (uf, - -iik-wwx)) s.t. wfA wook + L. (Xe
Emurch()) =
(15)
where WIA = Na/Now is the fraction of the population in the kth target distribution.
Mass Conservation Loss all times E [0, 1], we define a mass loss Luuss that enforces the sum of the weights of all + branches matches the true total weight at time / denoted as w total
To ensure that the growth rate satisfies conservation total mass at
Lmaas (0) = [[. Bip. 1
(
(
K W
tntal we
K
max (0. -0:,, de
(16)
where max(O, - w assigns an additional linear penalty for negative weight predictions. For the balanced branched SBM problem where the total mass is conserved, we have no!otal
5

---

Published as a conference paper at ICLR 2026
Training the Growth Networks In Stage 3, we train the growth networks { k by fixing the weights the flow networks {u[, 18 and minimizing the weighted combined loss growwhh with an additional growth penalty term 1 ensure coercivity of Cgrowch
Cgrewth (Ø) = Aenorgy Lonery (0, 6) + AmutchCmnnhh + Amsss Cmass (() + Agrowth
1
(17)
We show in Lemma C.2 that the optimal growth rates across the K secondary branches are non- decreasing; however, mass destruction can still be modeled by defining an additional branch with target weight equal to the ratio of mass lost over / [0, 1]. To ensure that the of optimal growth functions g" == 1 exists, we establish Proposition 4.2 (see proof in Appendix C.4).
Proposition 4.2 (Existence Optimal Growth Functions). Assume the state space X S R a bounded domain within Rd. Let the optimal probability density branch be known non-negative function bounded [0, 1), denoted as Pi,, x [0, 1] - [0, (0,1]= L X x [o, IJ).
By Lemma C.2, we can define the set of feasible growth functions the set of square-integrabbe functions 1.2 as
:= (g (gh.0.... 1t., (t,KK)
E (2 [0, 1]:1 1kk 2 1, (, 2 0, Σ k= U1A 0} (18)
Let the growth loss be the functional C(g)= (2 x 0. 0.11:KK+1 - R. Then, there exists an optimal function g" = (g....... g.,x) x) E L x [0, [0,1]: where g" E 9 such that
L(g") = gns ((g) which can be obtained by minimizing L(g) over g.
Final Joint Training In the final Stage 4, train the weights for both the flow and growth networks {ul Ie K by minimizing Sgown from Stage in addition to reconstruction loss frecoms that ensures the endpoint distribution at time i = 1 is maintained.
Crecom (0)
Bpl
Σ
max 0. (0, 21 112 121l--)) 55]
(19)
(#1,.,
where N₁, (#1,.) .. is the set of n-nearest neighbors to the reconstructed state i1 ~ PI,k at time = from the data points 11,A ~ TI.A at time /. = 1.
Our multi-stage training scheme decomposes the Branched CondSOC problem into two parts. We first independently learn optimal drift field for each branch, which is vector field over the state space that propagates mass flow in the direction of each target distribution. Then, we fix the drift fields and learn the growth dynamics that determine the optimal distribution of mass over the branches.
EXPERIMENTS
We evaluate BranchSBM on a variety of branched matching tasks with different state costs Vı(X,) including multi-path LiDAR navigation (Section 5.)), modeling differentiating single-cell population dynamics (Section 5.2), and predicting heterogeneous cell-states after perturbation (Section 5.3).
5.1 BRANCHED LIDAR SURFACE NAVIGATION
First, we evaluate BranchSBM for navigating branched paths along the surface of a 3-dimensional LiDAR manifold, from initial distribution to two distinct target distributions (Figure 3).
Setup We define a single initial Gaussian mixture To and two target Gaussian mixtures T1,0+ Æ1,1 on either side of the mountain (Figure 3). We sample 5000 points i..d. from each of the Gaussian mixtures and assign all endpoints a target weight of w w = 0.5. To ensure trajectories follow the LiDAR manifold, metric (Kapušniak et al., 2024; Arvanitidis et al., 2016), which assigns lower costs in regions near coordinates in the LiDAR dataset. Further experimental details provided in Appendix E.3.
define the state cost V.LAND
the data-dependent LAND

---

Published as a conference paper at ICLR 2026
Energy
Table 1: Benchmark of BranchSBM against single-branch SBM on multi-path surface navigation. Wasserstein distances (Wi and W2) between the reconstructed and ground truth distributions with Neteps 100 Euler steps at time = from validation samples in the initial distribution. Results averaged over indeppndennt runs.
Time
Time
Figure 2: Plot weight (left) and energy (right) calculated with (14) cach branch over time t E [0, 1]. Mass is transferred from the primary branch to branch I, and both converge to the target weight of 0.5 at Both plots represent the average trajectories from samples in the validation set.
Model
M4 (
Single Branch SBM
0.878 cm
BranchSBM
0,239
0.309
and Targat Datr butlons To and [1) Flu
Learned Dranching Dritx (edg (x))
Branch (x))
Figure 3: Application of BranchSBM Learning Branched Paths LiDAR Manifold. Plots of the initial and target distributions, learned interpolants, learned branched trajectories on the LiDAR manifold.
Results We show that BranchSBM can learn distinct, non-linear branched paths that curve along the 3-dimensional LiDAR manifold while minimizing the kinetic energy and state-cost. From the mass and energy curves in Figure 2, gradually transferred to the secondary branch (branch 1) over t E [0, 1]. with both curves converging 1. As mass is transferred, the slope of the cumulative energy curve
that begins in the primary branch (branch 0) and is
the target weight of 0.5 decreases in branch 0 and increases in branch I, reflecting the true energy dynamics. In Figure 3, we observe that the branching occurs at the edge of the inclined mountain, indicating that the model can determine the optimal branching time based on the paths of lowest potential energy. As shown in Table 1, BranchSBM reconstructs the endpoint distributions with significantly higher accuracy in comparison to single-branch SBM. In total, we demonstrate the capability of BranchSBM to learn branched trajectories on complex 3D manifolds.
5.2 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. Here, we demonstrate this capability mouse hematopoiesis data (Sha al., 2023; Weinreb et al., 2020) and pancreatic 3-cell differentiation data (Veres al., 2019).
Mouse- Hematopoomoss Results We use a dataset consisting of mouse hematopoiesis scRNA-seq data analyzed by lineage tracing technique from (Sha et al., 2023; Weinreb al., 2020). This data contains three time points t; for i {0, 1,2] that are projected to two- wo-diiensioona representations 2 referred to as force-directed layouts or SPRING plots. From the plotted data, we can observe two clear branches that indicate the differentiation of progenitor cells into two distinct cell fates (Figure 4). We use k-means clustering to define two distinct target distributions IT,00 and I,, of samples at time &2 and set their target weights equal to W1.0 W1,1 0.5 due to the equal ratio of cells (Figure 10). We used samples across all time steps t, for E {0, 1, to define the data manifold via the LAND metrie VIAND BranchSBM was trained pairs sampled only from and t2, and samples from t were held out for evaluation. For comparison, we trained a single- branch
SBM model with both clusters at lo the target distribution.
7

---

Published as a conference paper at ICLR 2026
Single Branch
Simulated at timatt
Simuliseed State
Leamed Trajetorien
Figure 4: Application of BranchSBM on Modeling Differentiating Single-Cell Population Dynamics. Mouse hematopoiesis scRNA-seq data is provided for three time points lo, (1, 42. (A) Simulated states (top) and trajectories (bottom) at time t1 using single-brand SBM. (B) Simulated states with BranchSBM at fi (t = 0.5) and (C) tz (t 1). (D) Leamed trajectories over the interval t [fo, tu] on validation samples.
Table 2: Results for pancreatic 3-cell differentiation experiment. I-Wasserstein distancc (W1) of intermediate distributions generated by DcepRUOT (Zhang ct al., 2025c), CytoBridge (Zhang et al., 2025d), and BranchSBM (Ours) at eight different time points. i denotes values taken from (Zhang et al., 2025d).
Model
Wi
W;
Wi
W
W2
W;
W2
DeepRUOT
8 8.0447 10.000
8.0773 0011
7,6301 go00
8. g0044amm
7.90181oo11
8. 8,3977 10.013
7,8346
CytoRridge 1
2.044820.00
8.0771 - 0021
7.629920.002
6.066620004
7.9018400117
8.397440am02
7.8340)
BranchSBM (Ours)
11.9774000000
7-4648
noi
11.620450000
11. 2503 2008
10.2888_6 0050
8,730) 20.0030
6. 6.8702g6 00c0
After simulating validation samples from the initial distribution Do ~ To. we evaluate the recon- structed distributions at the intermediate held-out time point t and final time point 12 (t 1). In Figure 4A, we observe that single-branch SBM trained with a single target distribution pl con- taining both terminal fates fails to learn distinct branched trajectories, and the simulated cell states at time t2 do not reach either of the terminal distributions. In contrast, we show that BranchSBM simulates branched states at intermediate time steps not included in the training data while accurately reconstructing both target distributions with significantly lower 1-Wasserstein and Wasserstein distances compared to the single-branch SBM model (Figure 4B-D; Table 3).
Table 3: Results for Modeling Single-Cell Differentiation. Wasser- stein distances (W2 and Wa) between simulated and ground-truth cell distributions at time i and f2 on the validation dataset. BranchSBM reconstructs both intermediate and terminal states significantly better than single-branch SBM. Results are averaged over 5 independent runs.
Pancreatic 3-Cell Differentia- tion Results We use a dataset consisting of pancreatic ß-cell differentiation data (Veres et al., 2019) containing 51, 274 cells col- lected over eight time points as they evolve from human pluripo- tent stem cells to pancreatic B-
Model
Single Branch SBM
BranchSBM
Time +
Wi (4)
W2 (L)
Wi (L)
Wa (L)
ti
0.58220.020
0. 0.7032...m 008
0.36620.080
0. 474
like cells. This data contains eight time points 4 for i E (0, (O.....7] 7} that are projected to 30- dimensional PC representations
l2
0.94022.077
1.037
0.210
0. 26.
E 30 We use Leiden clustering to define K = 11 distinet target distributions of samples at time 17 and set their target weights relative to the total weight at tp (App E.4). We used samples across all time steps to define the data manifold via the RBF metric R/B BranchSBM was trained on pairs sampled only from to and ł7, and the distributions at all intermediate time points were inferred from learning to minimize the distance from the data manifold over time. For baselines, we compared two SOTA methods for single-cell trajectory modeling: DeepRUOT (Zhang et al., 2025c) and CytoBridge (Zhang et al., 2025d), which are trained explicitly on intermediate snapshots to reconstruct the stochastic dynamics of cells that evolve independently.
8

---

Published as a conference paper at ICLR 2026
Notably, BranchSBM not only reconstructs the multi- terminal distribution at t7 with superior accuracy against all baselines, but also produces intermediate trajectories that are competitive with models trained directly on intermediate snapshots using explicit reconstruction losses (Table 2; Appendix Figure 9). Leveraging the RBF state cost V/RBF, which encourages trajectories to remain on the underlying data manifold, BranchSBM effectively captures the true differentiation dynamics through the combined influence of the neural interpolant and the path-energy objective. These results demonstrate that BranchSBM scales reliably to a large number of branches and is particularly advantageous in settings where intermediate timepoints are sparse or unavailable, allowing the model infer biologically meaningful trajectories even with limited temporal supervision.
f,y
5.3 CELL-STATE PERTURBATION MODELING
Predicting the effects of perturbation on cell state dynamics is crucial problem for therapeutic design. In this experiment, we leverage BranchSBM to model the trajectories of single cell line from a single homogeneous state to multiple heterogeneous states after a drug-induced perturbation. We demonstrate that BranchSBM is capable of modeling high-dimensional gene expression data and learning branched trajectories that accurately reconstruct diverging perturbed cell populations.
Setup For this experiment, we extract the data for single cell line (A-549) under perturbation with Clonidine and Trametinib at 5 ,dl, selected based on cell abundance and response diver- sity from the Tahoe-100M dataset (Zhang et al., 2025a). Since both drug perturbation datasets contained over 60K genes, selected the top 2000 highly variable genes (HVGs) based on nor- malized expression and performed principal component analysis (PCA) to find the top PCs that capture the variance in the data. We set the initial distribution t 0 to be control DMSO- treated cell population and the target distributions at = to be distinct clusters in the drug- treated cell population. After clustering, we identified two divergent clusters in the Clonidine- perturbed population and three in the Trametinib-perturber population (Appendix Figure 11).
To determine the weights of each branch, we take the ratio of each clus- ter with respect the total perturbed cell population (Appendix Table 9). For both experiments, we simulated the top 50 PCs, which capture approx- imately 38% of the variance the dataset. To further evaluate the scal- ability of BranchSBM on simulating trajectories in high-dimen nigh-dimensional state spaces, we simulated the top 100 and 150 PCs for Clonidine and compared the performance across dimensions.
Table 4: Results for Clonidine Perturbation Modeling for In- creasing Principal Component Dimensions. Maximum mean discrepancy (MMD) across all PCs and Wasserstein distances (Wi and W22 of top 2 PCs between ground truth and reconstructed distri- 1 simulated from the validation data O. Results for single-branch SBM (50 PCs) and BranchSBM (2 branches) were averaged over 5 independent runs.
butions
Model
RBF-MMD (1)
Wi (L)
W2 (L)
Single Branch SBM (50 PCs)
0.27910.024
5.12110.509
6.14910.463
BranchSBM 50 PCs 100 PCs 150 PCs
Given that the intermediate trajectory between the control and perturbed state is unknown, we assume that the optimal trajectory both minimizes the kinetic energy of the drift field while minimizing the distance from the space of feasible cell states. We define the state cost Vi(X,) with the RBF metric (Kapuániak et al., 2024; Arvanitidis et al., 2016), which pushes the intermediate trajectory to lie near states represented in the dataset. Further details are provided in Appendix E.5,
0.065
1.076
1.224
0.053
1.832
2,037 10.174
0.083
1.72210.064
1.931 10.035
After multi-stage training of BranchSBM with d E (50, 100, 150} PCs and two branched endpoints (Appendix Figure 6A), we simulated the final perturbed state of each branch at time = 1 from the samples in the initial validation data distribution 2g ~ no corresponding to the control DMSO condition. In Appendix Figure 6C, we demonstrate that BranchSBM accurately reconstructs the ground-truth distributions of endpoint 0 (top row) and endpoint I (bottom row) across increasing PC dimensions, capturing the location and spread of the dataset. To prove the necessity of our branched framework, simulate the target distribution with only single endpoint distribution Pl containing both clusters with single-brandd SBM and show that it only reconstructs the population of cells in endpoint O, which represent cells closest to the control cells along PC2, and
Clonidine Perturbation Results

---

Published as a conference paper at ICLR 2026
Trmmtti Perturtation Data
Branch
Energy Dwildtinn Sranch
Maas Deokition Branch
tar +11
Figure 5: Results for Trametinib Perturbation Modeling with BranchSBM. (A) Gene expression data of DMSO control (f = 0) and cells after treatment with M Trametinib (t = 0) with three distinct endpoints (purple, tuquuoisc, and pink). (B) The simulated endpoints of the top 50 PCs at f = 1 on the validation data for each branch. (C) The evolution of cumulative energy across & [0, 1] calculated (14) along each branched trajectory after Stage (growth with fixed drift) and Stage (joint) training. (D) The evolution mass across t E [0, 1] along each branched trajectory with target weights of w1.a = 0.603, W1,1 0.255 and w1.2 = 0.142.
fails to differentiate cells in cluster that differ from cluster 0 in higher-dimensional PCs (Appendix Figure 6B). Concretely, BranchSBM used across all PC dimensions outperforms single-branch SBM on only 50 PCs (Table 4), indicating that BranchSBM is required to model complex perturbation effects in high-dimensional gene expression spaces.
Trametinib Perturbation Results We further show that BranchSBM can scale beyond two branches, by modeling the perturbed cell population of Trametinib-treated cells, which di- verge into three distinct clusters (Figure 5A). We trained BranchSBM with three endpoints and single- branch SBM with one endpoint containing all three clusters the top 50 PCs.
After simulating the trajectories time t E [0, 1] the vali- dation cells in the control popu- lation, we show that BranchSBM generates clear trajectories to all three branched endpoints (Figure 5B) and reconstructs the over- all target distribution with lower error compared to single- branch SBM (Table 5). In Figure SC and D, we plot the evolution of cumulative energy calculated Levery (0, σ) (14) and weight of each branch over t [0, 1], demonstrating that BranchSBM's multi- stage training scheme effectively learns the optimal trade-off between minimizing the energy across trajectories and matching the target weights of each branch.
Table 5: Results for Trametinib Perturbation Modeling, Maximum- mean discrepancy (MMD) across all 50 PCs and Wasserstein distances (Wi and W2) of top 2 PCs between ground truth and reconstructed distributions at simulated from the validation data at = 0. Results were averaged over 5 independent runs.
Model
RBF-MMD (4)
Wi
W2 (4)
Single Branch SBM BranchSBM
0. 0.246
5.428
6.426 186
0.053
0.83840.061
0.973
CONCLUSION
In this work, we introduce Branched Schrödinger Bridge Matching (BranchSBM), a novel matching framework that solves the Generalized Schrödinger Bridge (GSB) problem from initial distribution to multiple weighted target distributions through the division of mass across learned branched trajectories. BranchSBM solves the branched Schrödinger bridge problem defined as of Unbalanced Conditional Stochastic Optimal Control tasks, parameterizing branch-specific velocity and growth rates with neural networks predict system trajectories with single inference simulation. Through applications to nonlinear 3D navigation, cell differeetiiaionn, and perturbation-inde expression, we demonstrate that BranchSBM provides a unified and flexible framework for modeling complex branched dynamics across biological and physical systems.
gene
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
We have taken multiple steps to ensure the reproducibility of BranchSBM. All theoretical derivations are fully detailed in the main text and Appendix, with proofs Propositions 3.1-4.2 provided in Sections C.I-C.4. Our multi-stage training procedure is described explicitly, and pseudocode is given Algorithm I. Experimental setups for each task are described in detail, including synthetic LiDAR navigation (Appendix E.3), differentiating single-cell dynamics (Appendix E.4), and cell- state perturbation modeling (Appendix E.5). We specify all hyperparameters and architectures in Table 11, evaluation metries such Wasserstein distances and MMD are defined in Equations 102, 103, and 104, and additional loss functions are given Equations 12-17. The codebase is freely accessible the academic community https: com/ sopht ang/ Bengg//rannn BranchSBM and httttsss//nggggggfaeeeeeee// /ataaaaeeeeetttttteeeeaaaa///a/e
ETHICS STATEMENT
This work focuses on the development theoretical and computational framework for learning branched Schrödinger bridges in generative modeling. All experiments are conducted on publicly available or synthetic datasets, including 3D LiDAR manifolds, mouse hematopoiesis scRNA-see data, and the Mahoe-100M perturbation dataset, and do not involve any personally identifiable, sensitive human subject data, or animal experiments. BranchSBM provides a general-purpose algorithm for modeling stochastic processes with branching dynamics, with potential applications in biological modeling, navigation, and physics. Possible risks, such as misuse in simulating sensitive biomedical data, are mitigated by the fact that BranchSBM is general generative framework that requires curated training datasets and does not directly generate actionable biological interventions. We believe the potential benefits of this work, in advancing generative modeling theory and enabling more accurate modeling of branching systems in biology and beyond, outweigh possible risks.
DECLARATIONS
Acknowledgments We thank Mark Systems, Penn Computing and Educational Technology Services (CETS), and Penn Advanced Research Computing Center (PARCC), for providing database and hardware support that has contributed to the research reported within this manuseript.
Author Contributions S.T. devised and developed model architectures and theoretical formulations, and trained and benchmarked models. Y.Z.. advised model design and theoretical framework, and processed data for training. S.T. drafted the manuscript and S.T. and Y.Z. designed the figures. A.T. reviewed mathematical formulations and provided advising. P.C. designed, supervised, and directed the study, and reviewed and finalized the manuscript.
Data and Materials Availability munity httpst://ggtteeee
The codebase is freely accessible to the academic com- com/ sophtanggBBaacchssM BranchsBM and https: httttp::/hugggingfaees co/Chatterjeelabb/ Braaeehhaabbaaaa........
Funding Statement
This research was supported by NIH grant R35GM155282 to the lab of P.C.
Competing Interests P.C. is a co-founder of Gameto, Inc., UbiquiTxx, Inc., AtomBioworks. Inc., and Recognition Bio, Inc., and advises companies involved biologies development and cell engineering. P.C.'s interests reviewed and managed by the University of Pennsylvania in accordance with their conffict-of-interes policies. S.T., Y.Z., and A.T. have no conflicts of interest to declare.
REFERENCES
Sayali Anil Alatkar and Daifeng Wang. Artemis integrates autoencoders and schrodinger bridges to predict continuous dynamics of gene expression, cell population and perturbation from time-series single-cell data. bioRxiv, pp. 2025-01, 2025.
11

---

Published as a conference paper at ICLR 2026
Georgios Arvanitidis, Lars Kai Hansen, and Søren Hauberg. A locally adaptive normal distribution.
Advances in Neural Information Processing Systems, 2016.
Lazar Atanackovic, Xi Zhang, Brandon Amos, Mathieu Blanchette, Leo J Lee, Yoshua Bengio, Alexander Tong, and Kirill Neklyudov. Meta flow matching: Integrating vector fields on the wasserstein manifold. International Conference Leurning Representations, 2024.
Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 2021.
Aymeric Baradat and Hugo Lavenant. Regularized unbalanced optimal transport as entropy mini- mization with respect branching brownian motion. arXiv preprint arXiv:2111.01666, 2021.
Bernard Beauzamy, Introduction to Banach spaces and their geometry, volume 68. Elsevier, 201 1.
Barbora Benešová and Martin Kružík. Weak lower semicontinuity of integral functionals and applications. arXiv preprint 1rXiv:1601.00390 2016.
Charlotte Bunne, Ya-Ping Hsich, Marco Cuturi, and Andreas Krause. The schrödinger bridge between gaussian measures has closed form. International Conference Artificial Intelligence and Statistics, 2022a.
Charlotte Bunne, Laetitia Papaxanthos, Andreas Krause, and Marco Cuturi. Proximal optimal transport modeling of population dynamics. In International Conference on Artificial Intelligncee and Statistics, pp. 65 6511-6528. PMLR, 2022b.
Charlotte Bunne, Stefan G Stark, Gabriele Gut, Jacobo Sarabia Del Castillo, Mitch Levesque, Kjong- Van Lehmann, Lucas Pelkmans, Andreas Krause, and Gunnar Rätsch. Learning single-cell perturbation responses using neural optimal transport. Nature methods, 20(11):1759-1768 2023.
Tianrong Chen, Guan-Horng Liu, and Evangelos A. Theodorou. Likelihood training of schrödinger bridge using forward- bowwardbackwarrd sdes theory. International Conference on Learning Representations, 2021a.
Tianrong Chen, Guan-Homg Liu, Molei Tao, and Evangelos A Theodorou. Deep momentum multi-marginal schrödinger bridge. 37th Conference on Neural Information Processing Systems, 2023.
Yongxin Chen and Tryphon Georgiou. Stochastic bridges of linear systems. IEEE Transactions on Automatic Control, 6125:5266-33, 2016. doi: 10. 10.1109/AA.22115244055
Yongxin Chen, Tryphon Georgiou, and Michele Pavon. On the relation between optimal transport and schrödinger bridges: A stochastic control viewpoint. Journal of Optimization Theory and Applications, 2016.
Yongxin Chen, Tryphon T. Georgiou, and Michele Pavon. The most likely evolution of diffus- ing and vanishing particles: Schrodinger bridges with unbalanced marginals. arXiv preprint arXiv:2108.02889, 2021b.
Lénaïc Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xaii Xavier Vialard. Unbalanced optimal transport: Dynamic and kantorovich formulations. Journal of Functional Analysis, 274(11): 3090-3123, 2018. ISSN 0022-1236.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transportation distances. Advances in Neural Information Processing Systems, 2013.
Bernard Dacorogna. Direct Methods in the Calculus of Variations. Springer New York, 1989. ISBN 9780387357799.
Valentin De Bortoli, James Thornton, Jeremy Heng, and Amnad Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances Neural Information Processing Systems, 2021.
12

---

Published as a conference paper at ICLR 2026
Valentin De Bortoli, Iryna Korshunova, Andriy Mnih, and Arnaud Doucet. Schrodinger bridge flow for umpaired data translation. Advances in Neural Information Processing Systems, 37: 103384-103441, 2024.
Atray Dixit, Oren Parnas, Biyu Li, Jenny Chen, Charles P Fulco, Livnat Jerby-Arnon, Nemanja D Marjanovic, Danielle Dionne, Tyler Burks, Raktima Raychowdhury, et al. Perturb-seq: dissecting molecular circuits with scalable single-cell ma profiling of pooled genetic screens, cell, 167(7): 1853-1866, 2016.
Alice Driessen, Benedek Harsanyi, Marianna Rapsomaniki, and Jannis Born. Towards generalizable single-cell perturbation modeling via the conditional monge gap. arXiv preprint XXv:22544008326 2025.
Jhanvi Garg, Xianyang Zhang, and Quan Zhou. Soft-constrained schrödinger bridge: a stochastic control approach. In International Conference on Artificial Intelligence and Statistics, pp. 4429- 4437. PMLR, 2024.
George I Gavriilidis, Vasileios Vasileiou, Aspasia Orfanou, Naveed Ishaque, and Fotis Psomopoulos. A mini-review on perturbation modelling across single-cell omic modalities. Computational and Structural Biotechnology Journal, 23: 886, 2024.
Nikita Gushchin, Sergei Kholkin, Evgeny Burnaev, and Alexander Korotin. Light and optimal schrödinger bridge matching. In Forty-first International Conference on Machine Learning, 2024.
Guillaume Huguet, Daniel Sumner Magruder, Alexander Tong, Oluwadamilola Fasina, Manik Kuchroo, Guy Wolf, and Smita Krishnaswamy. Manifold interpolating oraimoltrransport flows for trajectory inference. Advances in neural information processing systems, 35:29705-2971 88, 2022.
Kacper Kapušniak, Peter Potaptchik, Teodora Reu, Leo Zhang, Alexander Tong, Michael Bronstein, Avishek Joey Bose, and Francesco Di Giovanni. Metric flow matching for smooth interpolations on the data manifold. Advances in Neural Information Processing Systems, 2024.
Jun Hyeong Kim, Seonghwan Kim, Seokhyun Moon, Hyeongwoo Kim, Jeheon Woo, and Woo Youn Kim. Discrete diffusion schrödinger bridge matching for graph transformation. International Conference on Learning Representations, 2024.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International
Conference for Learning Representations, 2014.
Dominik Klein, Théo Uscidda, Fabian Theis, and Marco Cuturi. Genot: Entropic (gromov) wasser- stein flow matching with applications to single-cell genomics. Advances in Neural Information Processing Systems, 37:103897-103944, 2024.
Mutsumi Kobayashi, Misato Kobayashi, Junko Odajima, Keiko Shioda, Young Sun Hwang, Kotaro Sasaki, Pranam Chatterjee, Christian Kramme, Richie E Kohman, George M Church, al. Expanding homogeneous culture of human primordial germ cell-like cells maintaining germline features without serum or feeder layers. Stem Cell Reports, 17(3):07-521, 2022.
Alexander Korotin, Nikita Gushchin, and Evgeny Burnaev. Light schrödinger bridge. arXiv preprint arXiv:2310.01 174, 2023.
Christian Kramme, Alexandru M. Plesa, Helen H. Wang, Bennett Wolf, Merrick Pierson Smela, Xiaoge Guo, Richie Kohman, Pranam Chatterjee, and George M. Church. An integrated pipeline for mammalian genetic screening. Cell Reports Methods, 1(6): October 2021.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. International Conference on Learning Representations, 2023.
Peter Lippmann, Enrique Fita Sanmartín, and Fred A Hamprecht. Theory and approximate solvers for branched optimal transport with multiple sources. Advances in Neural Information Processing Systems, 25:267-279, 2022.
Guan-Horng Liu, Tianrong Chen, Oswin So, and Evangelos A. Theodorou. Deep generalized schrödinger bridge. Advances in Neural Information Processing Systems, 2022.
13

---

Published as a conference paper at ICLR 2026
Guan-Horng Liu, Yaron Lipman, Maximilian Nickel, Brian Karrer, Evangelos A. Theodorou, and Ricky T. Q. Chen. Generalized schrödinger bridge matching, International Conference on Learning Representations, 2023a.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. International Conference on Learning Representations, 2023b.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. International Conference for Learning Representations, 2017.
Christian Léonard. A survey of the schrödinger problem and some of its connections with optimal transport. Discrete & Continuous Dynamical Systems-A, 2014.
Frederike Lübeck, Charlotte Bunne, Gabriele Gut, Jacobo Sarabia del Castillo, Lucas Pelkmans, and David Alvarez-Melis. Neural unbalanced optimal transport via cycle-consistent semi-couplings. arXiv preprint arXiv:2209. 15221, 2022.
Stathis Megas, Daniel G Chen, Krzysztof Polanski, Moshe Eliasof, Carola-Bibiane Schönlieb, and Sarah A Teichmann. Estimation of single-cell and tissue perturbation effect in spatial transcrip- tomics via spatial causal disentanglement. In The Thirteenth International Conference on Learning Representations, 2025.
Kirill Neklyudov, Rob Brekelmans, Alexander Tong, Lazar Atanackovic, Qiang Liu, and Alireza Makhzani. A computational framework for solving wasserstein lagrangian flows. International Conference on Machine Learning, 2023.
Maxence Noble, Valentin De Bortoli, Arnaud Doucet, and Alain Durmus. Tree-based diffusion schrödinger bridge with applications to wasserstein barycenters. Advances in Neural Information Processing Systems, 36:55193-55236 2023.
Matteo Pariset, Ya-Ping Hsieh, Charlotte Bunne, Andreas Krause, and Valentin De Bortoli. Unbal- anced diffusion schrödinger bridge. arXiv preprint UrXiv:2306.09099, 2023.
Michele Pavon, Esteban G Tabak, and Giulio Trigila. The data-drive schroedinger bridge. arXiv preprint arXiv: 2018.
Michele Pavon, Giulio Trigila, and Esteban G Tabak. The data-driven schrödinger bridge. Communi- cations Pure and Applied Mathematics, 74(7):1545-1573 2021.
Stefano Peluchetti. Bm Coupled schrödinger bridge matching. arXiv preprint 2409:2440 2024.
Merrick Pierson Smela, Christian C Kramme, Patrick RJ Fortuna, Bennett Wolf, Shrey Goel, Jessica Adams, Carl Ma, Sergiy Velychko, Ursula Widocki, Venkata Srikar Kavirayuni, et al. Rapid human oogonia-like cell specification via transcription factor-directed differentiation. EMBO Reports, pp. 1-30, 2025.
L. C. G. Rogers and David Williams. Difffusions, Markov Processes and Martingales. Cambridge University Press, 2000. ISBN 978051180514
Martin Rohbeck, Charlotte Bunne, Edward De Brouwer, Jan-Christian Huetter, Anne Biton, Kelvin Chen, Aviy Regev, and Romain Lopez. Modeling complex system dynamics with flow matching across time and conditions. The Thirteenth International Conference on Learning Representations, 2025.
Yusuf Roohani, Kexin Huang, and Jure Leskovec. Predicting transcriptional outcomes of novel multigene perturbations with gears. Nature Biotechnology. 42(6):927-935, 2023.
Jayoung Ryu, Charlotte Bunne, Luca Pinello, Aviv Regev, and Romain Lopez. Cross-modality matching and prediction of perturbation responses with labeled wromo-wassersteerr optimal transport. International Conference on Artificial Intelligence and Statistics, 2024.
14

---

Published as a conference paper at ICLR 2026
Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal- analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, (76(4):928-944) 2019.
Erwin Schrödinger. Über die umkehrung der nalurgeselze. Verlag der Akademie der Wissenschaften in Kommission bei Walter De Gruyter u..., 1931.
Yutong Sha, Yuchi Qiu, Peijie Zhou, and Qing Nie. Reconstructing growth and dynamic trajectories from single-cell transcriptomics data. Nature Machine Intelligence, 6(1):25-39, 2023.
Ophir Shalem, Neville Sanjana, Ella Hartenian, Shi, David A Scott, Tarjei Mikkelsen, Dirk Heckl, Benjamin L Ebert, David E Root, John G Doench, al. Genome-scale crispr-cas9 knockout screening in human cells. Science, 343(6166):84-8)) 87, 2014.
Yunyi Shen, Renato Berlinghieri, and Tamara Broderick. Multi-marginal schrödinger bridges with iterative reference refinement. 28th International Conference on Artificial Intelligence and Statistics, 2024.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger 2023..
bridge matching. Advances in Neural Information Processing Systems, 2023.
Merrick D Pierson Smela, Christian C Kramme, Patrick RJ Fortuna, Jessica L. Adams, Rui Su, Edward Dong. Mutsumi Kobayashi, Garyk Brixi, Venkata Srikar Kavirayuni, Emma Tysinger, et al. Directed differentiation of human ipscs to functional ovarian granulosa-like cells via transcription factor overexpression. Elife, 12:e83291, 2023.
Vignesh Ram Somnath, Matteo Pariset, Ya-Ping Hsieh, Maria Rodriguez Martinez, Andreas Krause, and Charlotte Bunne. Aligned diffusion schrödinger bridges. Uncertainty in Artificial Intelligence, 2023.
J Michael Steele. The Cauchy-Schwarz master class: introduction to the art of mathematical inequalities. Cambridge University Press, 2004.
Panagiotis Theodoropoulos, Nikolaos Komianos, Vincent Pacelli, Guan-Horng Liu, and Evangelos A. Theodorou. Feedback schrödinger bridge matching. The Thirteenth International Conference on Learning Representations, 2024.
Panagiotis Theodoropoulos, Augustinos 1) Saravanos, Evangelos A Theodorou, and Guann-Horng Liu. Momentum mulli-margina schrödinger bridge matching, Conference on Neural Information Processing Systems, 2025.
Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In International conference on machine learning, pp. 9526-9536 PMLR, 2020.
Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 2024a.
Alexander Y. Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio. Simulation-free free Schrödinger bridges via score and flow matching. Proceedings of The 27th International Conference on Artificial Intelligence and Statistics, 238: 1279-1277 2024b.
Francisco Vargas, Pierre Thodoroff, Neil D. Lawrence, and Austen Lamacraft. Solving schrödinger bridges via maximum likelihood. Entropy, 2021.
Adrian Veres, Aubrey L Faust, Henry L Bushnell, Elise N Engquist, Jennifer Hyoje-Ryu Kenty, George Harb, Yeh-Chuin Poh, Elad Sintov, Mads Gürtler, Felicia W Pagliuca, et Charting cellular identity during human in vitro 3-cell differentiation. Nature, 569(7756)3368-333 373,, 2019.
Cédric Villani al. Optimal transport: old and new, volume 338. Springer, 2008.
15

---

Published as a conference paper at ICLR 2026
Dongyi Wang, Yuanwei Jiang, Zhenyi Zhang, Xiang Gu, Peijie Zhou, and Jian Sun. Joint velocity- growth flow matching for single-cell dynamics modeling. arXiv preprint arXiv:2505. 2025.
Gefei Wang, Yuling Jiao, Qian Xu, Yang Wang, and Can Yang. Deep generative learning via schrödinger bridge. International Conference on Machine Learning, 2021.
Caleb Weinreb, Alejo Forticluz-Fraaiielll,, Fernando D. Camargo, and Allon M. Klein. Lineage tracing on transcriptional landscapes links state fate during differentiation. Science, 367(6479), 2020.
Toshiaki Yachimura, Hanbo Wang, Yusuke Imoto, Momoko Yoshida, Sohei Tasaki, Yoji Kojima, Yukihiro Yabuta, Mitinori Saitou, and Yasuaki Hiraoka. scegot: single-cell trajectory inference framework based on entropic gaussian mixture optimal transport. BMC bioinformatics, 25(1):388,
2024.
Grace Hui Ting Yeo, Sachit D Saksena, and David K Gifford. Generative modeling of single- cell time series with prescient enables prediction of cell trajectories with interventions. Nature communications, 12(1):3222, 2021.
Jesse Zhang. Airol A Ubas, Richard de Borja, Valentine Svensson, Nicole Thomas, Neha Thakar, Ian Lai, Aidan Winters, Umair Khan, and Matthew G. et al. Jones. Tahoe- A giga-scale single-cell perturbation atlas for context-dependent gene function and cellular modeling. bioRxiv, 2025a.
Stephen Zhang, Anton Afanassiev, Laura Greenstreet, Tetsuya Matsumoto, and Geoffrey Schiebinger. Optimal transport analysis reveals trajectories in steady-state systems. PLoS computational biology, 17(12):e1009466, 2021.
Yuhui Zhang, Yuchang Su, Chenyu Wang, Tianhong Li, Zoe Wefers, Jeffrey Nirschl, James Burgess, Daisy Ding, Alejandro Lozano, Emma Lundberg, et al. Cellflow: Simulating cellular morphology changes via flow matching. International Conference Machine Learning, 2025b.
Zhenyi Zhang, Tiejun Li, and Peijie Zhou. Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport. International Conference on Learning Representations, 2025c.
Zhenyi Zhang, Zihan Wang, Yuhao Sun, Tiejun Li, and Peijie Zhou. Modeling cell dynamics and interactions with unbalanced mean field schrödinger bridge. The Thirty-Ninth Annual Conference on Neural Information Processing Systems, 2025d.
16

---

Published as a conference paper at ICLR 2026
OUTLINE OF APPENDIX
In Appendix A, provide extended background the relevant theory for learning optimal stochastic bridges (A.1) and simulating trajectories on the data manifold (A.2). In Appendix B, we discuss the relationship between our proposed formulation for BranchSBM and previous related works. Appendix C provides the theoretical basis for Sections and 4, including formal proofs for Proposition 3.1 (C.1), Proposition 3.2 (C.2), Proposition 4.1 (C.3), and Proposition 4.2 (C.4). We provide additional experiment results in Appendix D. In Appendix E, we describe further details on experiments and hyperparameters used including specific details for each experiment, including multi-path LiDAR navigation (E.3), modeling differentiating single-cells (E.4), and modeling cell- state perturbations (E.4). Finally, we provide the pseudo code for the multi-stage training algorithm in Appendix F.
Notation We denote the state space as X' C Rd and the time interval indexed with k E (O..... K}. We denote initial data distribution at time 1 = 0 as and the terminal as { The joint data distribution is denoted 10 and a pair of
[0, 1]. The branches are
data distributions at samples is given by (To. x1,k) ~ To For simplicity, we denote 1. = Λ ( L1,k dd = 1, (dlzo, 1x)..). Let Ut,x (Xxx,) denote the marginal velocity field, gr, denote the growth rate, and Pr.,(Xe.t denote the marginal probability density for the ktth branch, where we sometimes drop the input X2,x for simplicity. In addition, we denote the conditional velocity field and probability density as 1. Urjo,1.A (X æo, æ1.x) and P.jo..A: Pejo,1,x((x,, (xo F1, respectively. The optimal values for any quantity are superscripted with a symbol. We denote the parameterized flow
de neural networks as "4,K with parameters 0 and the growth neural networks with Dt.k with parameters
6. In the context of unbalanced endpoint distributions, we denote the true initial weight of a sample wn and the final weight of sample from the kth target distribution wi,x- The predicted weights generated from the growth dynamics are given by Wr, and we seek to match the predicted weight at time = 1 given by wi.k to the true weight w k- L denotes the space of square integrable functions and || I| be the -2-m in function space. L denotes the space of essentially bounded functions such that || < 00.
A EXTENDED THEORETICAL BACKGROUND
A.1 LEARNING OPTIMAL STOCHASTIC BRIDGES
Pinned-Don -Dow Stochastic Bridges Let Q E M be a Markovian reference path measure that evolves over E 10, 1] according to a drift field (x,)) : Rdd + Rd and stochastic d-dimensional Brownian motion B, E Rd. via the SDE
dx, f(X)) + o,d Xo To
(20)
Given Q, consider a stochastic process (X, 1 over the time interval t E [0, 1] pinned-down at the initial point Xg = 2o and final point X1 = x1 denoted as Q 0.1 Do. x1). Due to the endpoint conditions, Q 10 is not necessarily Markov, and evolves via the SDE
dx, (f((.)) + V. log (x)] + o,dB,, Xo ZDD
(21)
where Va log (1) Xx)) is mon-Maakovian score function that corrects the drift field f((X,) of the reference process such that it points toward the target endpoint Φ, Since log Quje(r ( is the log-likelihood that the final state satisfies the condition X, 11 the gradient defines how the log- loglliiei changes with respect to the changes in the state æ at time t. The drift moves æ in the direction of the largest increase likeliioddood given by the score function, which ensures that the process satisfies X1 æi following the theory of Doob's h-transform (Rogers & Williams, 2000). Now, we can define the conditional prohability distribution pe as mixture of pinned-down stochastic bridges over pairs of endpoints in the data coupling T no X m given by
/
(a():) = No,1 Orjo, (2: xo, x1)
Qrjo. (20/ xx To,1 (dao, dx,)
(22)
17

---

Published as a conference paper at ICLR 2026
2.. To simplify notation, denote each conditional bridge as Qrjo. (x, x1) Pejo.1 (X, æ1) and the joint distribution To,1 (æo, #1). Now, we can rewrite the marginal as
(x() - - I x To, (daoo, dxi) ) Bma.1 [oojo (X,
x0, xi)]
(23)
Furthermore, we denote Urjo,1 = u, 2o: x1) = (Xt + V log Q1je(211 Xx)) as the conditional drift that generates pt Pejo,1 (X, zo, x1). that satisfies the conditional Bekker-Planck equation
0 Ot Pejo1 - (PH0.) 2220222000020 o2
1
(24)
Definition A.1 (Reciprocal Class). Given our definition of the conditional bridge Q. can define the reciprocal class, denoted R(Q)), of the reference measure Q as the class of path measures that share the same bridge as Q defined as R(Q) {1I 11(Xx xi)) QXX+/ x.)).
Markovian Projections Given mixture of conditional stochastic bridges II Io, 1o under the reference measure Q that require knowledge of the joint distribution To.1 we aim to project to the space of Markovian measures M, where the drift dynamics (x,X,) is only dependent the current state X₁ and require no knowledge on the endpoints. This allows us to parameterize Markovian drift uf (X₂) that can transport samples from the initial distribution 2p ~ To to samples from the target distribution æ To do this, we define the Markovian projection of TI (Shi et al., 2023; Liu et al., 2023b).
Definition A.2 (Markovian Projection). Given a conditional bridge Q. that evolves via the SDE in (21). we define a Markovian projection of the mixture of bridges II = Ho, jo.1 Markov process MI* projm (11) E -M with the same marginals as 11 such that X, ~ II, for all t [0, 1}, X1 mi and evolves via the SDE
dX, (f((X.) + (X))dtt + o,dB, (x.))=oPBMm T Vv,, log (x)
(25)
(26)
where U1i is the conditional distribution of X1 under the mixture of bridges II and Nz, log Qij (X,x)) points in the direction of greatest increase in the log-likelihood of the target endpoint X1 under the reference process Q. In addition, the Markov projection M- minimizes the KL-divergence with the mixture of bridges MI arg minMEM KL [M) and can be obtained by parameterizing of (X.) and minimizing the dynamic formulation given by
MI( = == Pat Fag S.'
207 love. log Quje(mi/z.) of = (1=
(27)
In general, the Markovian projection reference measure Q does not preserve the conditional bridge and is not in the reciprocal class R(Q). The unique path measure IP that is the Markovian projection of Q is in the reciprocal class R(Q), and preserves the endpoint distributions, is called the Schrödinger Bridge.
Definition A.3 (Schrödinger Bridge). Given reference measure Q, initial distribution To- and final distribution ग A path measure P is the unique Schrödinger bridge if satisfies the Markovian projection Q such that P =rrj ( (Q).
is in the reciprocal class of Q i.e. IP E R(Q), such that it preserves the conditional bridge P(X,|æo,, 21) Q(X, Fo. 1
preserves the endpoint distributions Po no and P Wj.
The goal of Schrödinger Bridge Matching (SBM) is estimate the Schrödinger Bridge that transports samples from an initial distribution Ta to a final distribution म given the optimal reference dynamics. We further discuss previous approaches solving the SB problem in Appendix B.
18

---

Published as a conference paper at ICLR 2026
A.2 SIMULATING TRAJECTORIES ON THE DATA MANIFOLD
Riemannian Manifolds and Metrics Since the interpolant 2t,, learned in Stage 1 is defined by minimizing non-linear state cost (x(X.), the resulting velocity field (X,) lies on the tangent bundle T of a smooth d-dimensional manifold Ω E IR called a Riemannian manifold. Intuitively, be thought of smooth surface where the local curvature around point æ E Ω can be approximated by tangent space Tast that defines the set of directions in which can move along the manifold. These directions are defined by a tangent vectors u E TaS,, which push the point æ along the manifold.
Riemannian manifold
In Stage 2, we seek to parameterize vector field uf (X₂) that minimizes the angle from the tangent vector zt.n 0, at each point on the manifold optimized in Stage 1. To do this, we must define the concepts of length and angles on Riemannian manifolds. First, we define a location-dependent inner product in Riemannian manifolds known as the Riemannian metric gæ T x TaSt - IR that maps two vectors u, v E Tas2 to a scalar that describes the relative direction and length the two vectors. Formally, the Riemannian metric can be written the billinear and positive definite function
Yu + 0 gæ(u, u) > G
(u, = u G(x) = (u, Go) s.t.
(28)
which defines the norm of tangent vector as || = V (u, u). Now, we can decompose the Riemannian norm of the tangent vector || to get the loss defined in (12) as follows
Laj (1) = E E,, 001) mmm )
[l 11
(29)
(20, 001) )
(it. ()
(30)
(20,
[| 112 + (dt..: (G( -I).,))
(31)
Ec,(20.20) )
|| 112 + (IF..)))
(32)
where the state cost is defined as (a
= (dre,n) (G( T)dtt.g)-
Data-Depende State Cost Following Kapušniak et al. (2024), define the metric ma- trix (æ described previously as the data-dependent LAND and RBF metrics of the form GLAND( (z D) GRar(æ, D) (diag(h(z)) =I) which assigns higher cost (i.e. [[(())] is larger) when æ moves away from the support of the dataset D. Specifically, given dataset { j=1 we define the elements h1-AND (x:) E IRd that scales down each dimension j E {I...., d} in the LAND metric
D
N
2 | 2
LAND (x)
Σ (I) 2d)eexp i=1
(33)
202
where the exp term is positive when there is high concentration of data points around the point (i.e. ||æ æxll is small) and approaches 0 as the concentration of data around decreases (i.e. [x - æxl || is large). Writing (ir... G (z )ar..) in terms of (2:), we get
(1) = 1 w (der, G dit...)
(34)
(or, )2
= (x0,,ar J~
(35)
h ( ++ +
When hy (x) is large, the loss is minimized, and when (x) is small, the loss is scaled up. While the LAND metric effectively defines the data manifold in low dimensions, in high-dimensional state spaces, setting a suitable variance in h A.LAND( the data manifold without overfitting is challenging. To overcome this limitation, the RBF metric clusters the dataset into N₂ clusters with centroids denoted as t, E Rd and trains of parameters (Wa, N.
ensure that the path does not deviate far from
1 to enforce hy(#;; 1 for all points in the dataset such that points within the data
19

---

Published as a conference paper at ICLR 2026
manifold are also assigned hy (z)
1. Specifically, h3 is defined
N: Wn,j (2) exp
An.j |
RBF h
(x)
in 112
(36)
2
n=1
-2
An 2
x
In
12
(37)
C
DE C..
where On denotes the rith cluster, A is a tunable hyperparameter. and An.j is the bandwidth of cluster the jth dimension. To train the parameters, we minimize the following loss
CRBF((Wa,
Σ (1- 1 (z:))²
(38)
I,ED
In our experiments, we use the LAND metric for the LiDAR and mouse hematopoiesis datasets with dimensions d = 2 and d = 3 respectively, and the RBF metric for the perturbation modeling experiment with gene expression data of dimensions {50, 100, 150).
Sampling the Data Manifold In Riemannian geometry, each Euclidean step At-u (a:) following the tangent vector ui (x) along the manifold requires mapping back to the manifold via an exponential map XI+A1 expx ( At (z)). In general, computing the exponential map exp,, under the Riemamnian metric (zx) requires simulation the geodesic flow (i.e., approximating the final state ( = under the initial conditions 2o x. and dho un? (ax)).
While the manifold defined by the data-dependent metric Gæx, D) induces Riemannian geometry in Euclidean space that follows optimal cost structure, its underlying space is still Euclidean. Gav, D) just assigns varying costs of moving in the Euclidean space Rd Therefore, we can avoid computing the exponential map and generate trajectories with simple Euclidean Euler integration following
Tt 3t + At (xx))
(39)
where At = 1/Nasps is the discretized step size.
B COMPARISON TO EXISTING WORKS
In this section, we discuss the relationship between our proposed formulation for Branched Schrödinger Bridge Matching and related previous works. We establish reasons why BranchSBM is the theoretically optimal formulation to solve the problem of modeling stochastic dynamical systems with diverging trajectories over time by modeling branched Schrödinger bridges. We conclude by introducing an alternative perspective of the Branched GSB problem defined in Section 3.2 as the problem of modeling probabilistic trajectories of dynamic systems with nondeterministic states, which BranchSBM is well-suited to solve.
B. MODELING BRANCHED SCHRÔDINGER BRIDGES
Schrödinger Bridge Matching Computational methods for solving the Schrödinger Bridge (SB) problem for predicting trajectories between initial and target distributions have been extensively studied in existing literature (De Bortoli et al., 2021; Chen et al., 2021b; Korotin et al., 2023; Bunne
al., 2022a; Chizat et al., 2018; Liu et al., 2023a; 2022: Shi al., 2023; Kim al., 2024; Wang
et al., 2021; Tong et al., 2024b; Poluchetti, 2024; Bunne et al., 2022a; Somnath et al., 2023; Gushchin
al., 2024; Pavon et al., 2021; Garg et al., 2024; De Bortoli et al., 2024; Shen et al., 2024; Chen
et al., 2023; Noble et al., 2023). Previous work has framed the SB problem an entropy-regularized Optimal Control (EOT) problem (Cuturi, 2013; Léonard, 2014; Pavon et al., 2018) or a stochastic optimal control (SOC) problem (Chen et al., 2016; 2021a; Liu et al., 2023a), which we build in this work. To ensure that the intermediate trajectories follow a path distribution that approximates that observed in the data, multimarginal Schrödinger bridge matching ensures reconstruction of multiple intermediate marginal distributions at specified time steps along the path (Shen et al., 2024; Chen al., 2023; Theodoropoulos al., 2025).
20

---

Published as a conference paper at ICLR 2026
Stochastic Optimal Control Several works (Chen al., 2016; 202 Liu et al., 2023a) have
reframed the SB problem as a Stochastic Optimal Control (SOC) problem, which takes the canonical form
min [ Ppp [ (X.)ll VV(X.)) dt + Ep (x1))
(40)
We Jo
s.t. dx,- = ue( (x, + odBr Xo ~o
(41)
where (X1): X -> IR acts as a reconstruction loss that enforces that the distribution of X1 Pi matches the true distribution Ti. Due to the intractability of To, T1, GSBM (Liu et al., 2023a) uses spline optimization to learn an optimal Gaussian probability path Pžj0.1 (1,,, >LIa) using only samples 2g ~ To and æ1 ~ πi from the initial and terminal distributions. Although GSBM does not require knowledge of the densities, it follows an iterative optimization scheme that alternates between matching the drift uy given a fixed marginal pt, and updating the marginal given the drift. This strategy can get stuck in suboptimal solutions and is sensitive to the initialization of ug. GSBM is also limited to learning unimodal Gaussian paths between one source and one target distribution with balanced mass, making it unsuitable for modeling tasks with multimodal terminal distributions and splitting of mass over multiple distinct paths.
Regularized Unbalanced Optimal Transport Several previous works have studied the unbalanced optimal transport problem (Zhang et al., 2025c; Chen et al., 2021b; Lübeck et al., 2022; Pariset et al., 2023); however, these approaches address fundamentally different setting from the unbalanced Generalized Schrödinger Bridge (GSB) problem considered in this work. Specifically, DeepRUOT (Zhang et al., 2025c) solves the Regularized Unbalanced Optimal Transport (RUOT) problem by parameterizing the canonical stochastic bridge SDE dXx, = (dl + o, with probability flow ODE given by
dx, = {f((X))
1 Vo log (x,) dt
(42)
2
x (x,)
is the drift of the probability flow ODE that is learned along with the probability density (X,) to derive the drift of the SDE f(()).. Unlike GSBM (Liu et al., 2023a), which enforces hard constraint on the endpoints, DeepRUOT models a probability density flow by minimizing a reconstruction loss that encourages alignment with both intermediate and terminal distributions. While effective when intermediate distributions are observed, the method fails to learn meaningful trajectories in settings where these intermediate snapshots are unavailable, limiting its applicability in many real-world scenarios.
where (X,) X,)
Learning Diverging Trajectories with Single Target SBM To model diverging trajectories with single-branch SBM where the target distribution is multi-modal, we can follow a set of N samplees samples from the initial distribution no to determine the distribution of samples that end at each of the modes of the target distribution æ While this can estimate the splitting of mass across different trajectories, it does not explicitly learn the optimal distribution of mass in the latent space over time.
Furthermore, the path toward a specific cluster in the target distribution has lower potential than that of the other clusters, mode collapse could occur, where all samples follow the same trajectory without reaching the other clusters. BranchSBM learns to generate the correct mass distribution over each of the target states by optimizing the growth term with respect to the matching loss Lmatch defined in Equation 15.
With the standard SBM formulation, it is also challenging to determine the time at which branching occurs and the population diverges toward different targets, as all samples follow stochastic trajecto- ries. With BranchSBM, we model population-wide branching dynamics with growth networks that are trained to predict the origin of branch from having zero mass (w,, = 0) and the rate at which it grows/shrinks over time OrWi,k It. (Xt,k) The mass of a branch any given time step can be simulated with Week (X.,)) = Woo + Jo J
( (xs.))s..
Branching Dynamics While branched dynamics have been explored in the context of optimal transport (Lippmann al., 2022) and Brownian motion (Baradat & Lavenant, 2021), no previous
21

---

Published as a conference paper at ICLR 2026
work has explicitly formulated solved the branched Schrödinger bridge problem that seeks to match an initial distribution to multiple terminal distributions via stochastic bridges. For instance, the concept of Branching Brownian Motion (BBM) introduced in Baradat & Lavenant (2021) models population of particles that each independently follow stochastic trajectories according to standard Brownian motion with positive diffusivity v > 0. To model branching dynamics, each particle has an independent branching rate À > O that determines the probability that the particle undergoes branching event, defined as the particle dying and generating k new particles that then evolve independently. The number of generated particles is sampled from a probability distribution over non non-negative integers k ~ p = (PE) where piş is the probability of generating & particles the branching event and EKEN Ps 1. Given Ps. Qh ApR defines the rate of branching events that generate k new particles and defines new probability measure q = (qx) keN called the branching mechanism.
While BBM defines branching as the generation of additional particles from single particle following an independent, temporal probability measure q, this model fails to model the division of mass across multiple trajectories, where total mass remains constant but the mass of each branch changes. In addition, BBM assumes that each branched particle undergoes independent Brownian motion, without explicitly defining a terminal state or branch-specific drift. For these reasons, the BBM model is unsuitable for modeling branching in the majority of real-world contexts, such as cell state transitions, where undifferentiated cells split probabilistically into distinct fates rather than proliferating in number. In such systems, branching reflects a redistribution of probability mass over developmental trajectories, governed by underlying regulatory patterns rather than purely stochastic reproduction.
Where existing frameworks fall short is in modeling meaningful energy-aware, conditional stochas- tic trajectories with unbalanced and multi-modal dynamics, which we address in this work. Specifically, we formulate the Unbalanced CondSOC problem followed by the Branched CondSOC problem that defines set of optimal drifts and growth fields that define set of branched trajectories following optimal energy miriiiiiii trajectories defined by the state cost (.).. Instead of spline optimization, we leverage a parameterized network sP.., (20, x1) that learns to predict an optimal interpolating path given pair of endpoints.
To model dynamic growth of mass along branched trajectories, we initialize a normalized population weight of 1 at a primary branch (k = 0), that can split across K branches and generate weights {we, k== that evolve co-currently via learned growth rates { k This approach relies on learning the growth rate over an entire population of samples across all branches rather than learning independent growth rates of individual samples, which enforces stronger constraints during training to ensure that the model captures true population growth dynamics.
B.2 MODELING PERTURBATION RESPONSES
In single-cell transcriptomics, perturbations such as gene knockouts, transcription factor induction, or drug treatments frequently induce cell state transitions that diverge into distinct fates, reflecting differentiation, resistance, and off-target effects (Shalem et al., 2014; Kramme et al., 2021; Dixit al., 2016; Gavriilidis al., 2024; Zhang al., 2025a; Kobayashi al., 2022; Smela al., 2023; Pierson Smela et al., 2025; Yeo et al., 2021). Trajectory inference methods for modeling cell-state dynamics with single-cell RNA sequeneing data, including flow matching (Zhang et al., 2025b; Rohbeck et al., 2025; Atanackovic et al., 2024; Tong et al., 2024b; Wang et al., 2025), Schrödinger Bridge Matching (Alatkar & Wang, 2025; Tong et al., 2024b), and optimal transport (Zhang et al., 2025c; Tong et al., 2020; Bunne et al., 2023; Driessen et al., 2025; Huguet et al., 2022; Klein et al., 2024; Yachimura al., 2024; Schiebinger et al., 2019; Zhang et al., 2021; Bunne al., 2022b), have been widely explored. While many algorithms have been developed to predict perturbation responses (Bunne et al., 2023; Megas et al., 2025; Rohbeck et al., 2025; Roohani et al., 2023; Ryu et al., 2024), they either predict only the terminal perturbed state without modeling the intermediate cell-state transitions or model single unimodal perturbation trajectories. Schrödinger Bridge Matching frameworks like [SF] (Tong et al., 2024b) and DeepGSB (Liu al., 2022) have been shown to effectively model stochastic transitions in biological systems, offering better scalability and sampling efficiency than classical OT.
However, these models are still limited to modeling trajectories between a single pair of boundary distributions, limiting their ability represent divergent trajectories arising from the same perturbed
22

---

Published as a conference paper at ICLR 2026
initial state. This is key limitation when modeling processes like fate bifurcation pesttterturbaiion where a cell population exposed to the same stimulus may split into multiple phenotypically distinct outcomes. BranchSBM extends the SBM framework to support multiple terminal marginals, enabling modeling of stochastic bifurcations in a mathematically principled way. By learning a mixture of conditional stochastic processes from common source multiple target distributions, BranchSBM can capture the heterogeneity and uncertainty of cell fate decisions under perturbation. Moreover, it retains the empirical tractability of previous SB-based models, requiring only samples from distributions, and ensures that intermediate trajectories lie on the manifold of feasible cell states.
B.3 MODELING PROBABILISTIC TRAJECTORIES WITH BRANCHSBM
We conclude the discussion with alternative interpretation of the Branched Schrödinger Bridge problem that deviates from the branching population dynamics problem. We instead consider Branched SB as a probabilistic trajectory matching problem, where each branch one of multiple possible trajectories that sample X; could follow. Since single-path SBM learns only a single deterministic drift field u, (X,,) that determines the direction and flow of the SDE dXx, u (X,)de + odB,, the probabilistic aspect of the trajectory remains restricted to Brownian motion via odB, This fails to capture probabilistic dynamics where probability densities begin concentrated at a single state but evolve into multi-modal probability densities (Pt, k= over multiple different states, each of which evolve according to an SDE d'X,,,= = 4 (Xe, + adB, with an independent drift term ur (X),k) and noise scaling
Where other SBM frameworks fall short, BranchSBM is capable of modeling multiple probabilistic trajectories, where system begins single deterministic state X,- = zn with probability w1,0 1 and evolves via multiple probabilistic trajectories that diverge in the state space governed by the SDEs d'X,,k ur, (x,,(Xx + odB.. At time I, the system exists in non-deterministic superposition of states (Xt,e, Wr.k 1K k= each with a probability Wr such that Σ " = 1. In addition, BranchSBM can model the evolution of the probability weights w. by parameterizing the probabilisttic growth rates )K ) that preserve conservation of probability mass Σk= w = 1 at all times / [0, 1] by minimizing the mass loss (o)ss()) (16). This problem prevalent in many biological and physical systems, where system does not exist in single deterministic state, but rather superposition over a distribution of states.
Table 6: BranchSBM enables robust modeling of branched Schroödinger bridges compared to existing frameworks. BranchSBM can model both branching and unbalanced trajectories, follows intermediate trajecto- ries governed by a task-specific state cost, requires only endpoint samples for training, and samples trajectories from only single sample from the starting distribution.
Method
Models Branching
Models Unbalanced No
Intermediate Dynamics
Requirements for Training Samples from No,
Requirements for Sampling
Generalized SBM Liu al. (2023a)
No
Entropic OT with learned drift of
Endpoints (mo, p6.1 Sample
DeepRUOT (Zhang
Requires simulating multiple samples
Growth rate 9, (X.)
Regularized OT
Samples from Wo, Wi
al., 2025c)
with learned
æo do
drift of and density pf
BranchSBM (Ours)
Simulates divergent trajectories and terminal states from single sample ato
Branched growth rates (X.))
Branched drifts {14, a K that minimize state-cost
Samples from and clusters {m.
Sample
anp
23

---

Published as a conference paper at ICLR 2026
C THEORETICAL PROOFS
C.1 PROOF OF PROPOSITION 3.1
Proposition 3.1 (Unbalanced Conditional Stochastic Optimal Control). Suppose the marginal density can be decomposed as pe (Xx) = Pe(Xalao, 2 oo (dæo, dæi). where To,1 is a fixed joint coupling of the data. Then, we can identify the optimal drift un and growth gi that solves the Unbalanced GSB problem in (4) by minimizing the Unbalanced Conditional Stochastic Optimal Control objective given by
[." Eprion [() (( (00, x11 x1) +((()) we (X.) dt wū. wi
min (1mo:=1)=mo.1 Medde
(5)
S.f. dX. us (x 2s ddt adB., Xa = vo, X1
(6)
where We = w + So g (X,) ds is the time-dependert weight initialized at wh, ue the drift, gt the growth rate, and No,1 is the weighted coupling of paired endpoints (æo, wh: 2., w 20
Proof. We define the Unbalanced Generalized Schrödinger Bridge problem as the solution (u,, (t, 9r) to the energy minimization problem such that the unbalanced Fokker-Planck equation satisfied.
1
[l)a + vi(x.)
(wo+
+ L (x(X))ds)
min
Epp
(43)
Pt (x)
- (X.) -2 (x,(X,) + (X.)).x((
s.t.
(44)
Po No. P1=T T
Under the assumption that the joint probability To,1 x1) fixed over times E [0, 1] and the marginal probability can be factorized as pe (X,, Emo.. (Pejo, (X, we can decompose the minimization objective into
L.' Epp.
Me (X.)112 + v.(x.)
(wo + I (x()))da) dt
(45)
2
1/¹
100 1o0+ + L/' (xx)))d de (law total expectation)
Ero.
[ll(xx))1 + vi(X.)
2
[
Wa f. sx(X.)ds) de
= Epo.1
E
(x))( 3 vvx..))
(Fubini's theorem)
which can be solved by identifying the conditional drift u, that minimizes the expected objective value over all endpoint samples (æo, a1) To,1- Under similar assumptions, we can decompose all terms in the Fokker-Plance equation. For the left-hand side, we have:
0 Ht. / PX ai To. (dwo, dxy) Be )
0 ar Pu 0.1 (Xelwo, a01)
म. (X.)
(46)
For the divergence term, we have:
V (x.()))) - V. (m(x.)) / pejo. a (TO (dæo, day))
= (fx mi J (daro, dx1)
(47)
By Fubini's Theorem and the linearity of the divergence operator, we can switch the order of integration to get
V (xx()))((X))) V. V..( V.(u(X)))m (X1lx0, m1) iX Tu.1 (dxo, dx1))
B(rro 20 A (ue(Xe) (X,
(48)
24

---

Published as a conference paper at ICLR 2026
For the Laplacian term, have:
o²
o²
(X)(X
V.
(VVP((XX))
2
2
0²
V.
Pejo,1 (X2 2ro æ1 11)) )oo, (doo, dxx1)
2
d²
/
Vptoo (Xr xo. (20 xi)
To,1 (dxo, dx11)
2
σ²
16 (100,00)) |oTTT.
A (x,
(49)
2
For the growth term, we have:
ge (((()))((X)) = 9( / Prio, (Xelaox ai To, (daoo dry)
- f (x.())))oo... (Xel a1) T (dxo, dxa)
= -B J
ge (x(Xx))) (xe (X, m1)
(50)
Combining all the terms of the Fokker-Planck equation, we have shown that (44) can be rewritten as
θ कe Pejo.a (Xel (2o. mi) = V (X.]aj,. m1)) + (X. m1) + 9tPcjo.1 (X/woo m1)
(51)
Therefore, the Unbalanced GSB problem equivalent solving the Unbalanced CondSOC problem. and we conclude our proof of Proposition 3.1.
C.2 PROOF OF PROPOSITION 3.2
Proposition 3.2 (Branched Conditional Stochastic Optimal Control). For each branch, let (x,,(X.,,))== J Pt. (X (a, 21,k To (dæo, dær,k). )), where T0,1,k is the joint coupling dis- tribution of samples Do ~ from the initial distribution and 1..k from the kth target distribution. Then, we can identify the set of optimal drift and growth functions tut, Fi k=o that solve the Branched GSB problem (7) by minimizing the sum of Unbalanced CondSOC objectives given by
1K
L. EEriAAL
[ vi(X wro Jae
E(00:0100)~00.13 00>-000110 (
min K
B( S/' Ep.jo.1. [1 (X vi(X.,)) wei dl k=1
K
(8)
s.t. dx,,, ur.k Vr.. (x.
+ odB,, Xo = wo, X1,6 L1.8, wo,k 8c Wi. = wi.
(9)
where W100 1+ So 92,0 gx,0 (os is the weight of the primary paths initialized at 1 and
Wi,k So Ho,k (X.,k) (ds are the weights of the K secondary branches initialized at O.
Proof. We extend the proof of Proposition 3.1 to the branching case by defining each branch k an independent Unbalanced Generalized Schrödinger Bridge problem (4) given by
min ('' T.... | (X,,) Vo(Xe.)) wo, + S. go. (x, dt
(52)
Dy Pe - (x + Apra(Xea) + t. Po = no. PI,t = Ri,b Vk KK
s.t.
(53)
such that each branch independently solves the Bokker-Plance equation defined as Pr. (u..Ptt))) o" Apr.-- Now, show that the of unbalanced CondSOC problems still satisfies
25

---

Published as a conference paper at ICLR 2026
the global Fokker-Plance equation
a -V (WrPe) + d²
कe Pt
2 SPp
(54)
where we define as the weighted sum of the branched distributions given by
(x() Σ (et,
(55)
To obtain an expression for the global Pokker-Plande equation, we differentiate both sides and substitute the branched FPE follows
0 Pt
00
dt
L==
K
0
8 Week
01
Pt
Σ
Weik
H Perk
pe.
R=D
K
(-
o² Apr.
01
Pt
Whik
(utx)Ptx))
St.EPE
(substitute branched FPE and 0ewr. = gt.k)
K
0t
P1
Σ
We.K (uc,Pe,s) + Wr 2
Apr.k + Dt.bPe.k
k=0
o²
Pt
W. V. (ur,et)))) Σ
Wi,k
Apr
(56)
By linearity of the Laplacian, the diffusion term can be rewritten
K Wi.k
0² A
0²
A
Apt
(57)
2
2
K=
k=
Since the global growth term Σ 9t. &t.tPt is the sum of the growth dynamics across all branches and separate from the drift and diffusion dynamics, doesn't alter the direction or motion of particles along the branched fields. Thus, both the diffussin and growth terms satisfy the global FPE.
Now, set the divergence term in (56) equal the global divergence V (upr) and derive the expression for the total drift field ue (X,)) that satisfies the global FPE. By the linearity of the divergence operator, we get
K
V (upp:)
(-Wi V (u.,Pttk))
K
(u.p:)
Σ
1==0
(usp.) = V.
Wt, &P.
(58)
k=0
Under the global FPE constraint, the drift u, is defined as the wesghted average of the drift fields for each branch, given by ux(X,) pr(X.) Σk Wi, (Xe,x) u,, (X. Intuitively, this means that in the global context, for any X; æ., the drift of state X; along the dynamics of branch & is scaled by the weight of x; at time along branch and the ratio of probability density of æ, under branch & over the total probability density of æ; across all branches. Therefore, our definition of the weighted drift decoupled over individual branches satisfies the global FPE equation in (54), and this concludes the proof of Proposition 3.2.
26

---

Published as a conference paper at ICLR 2026
Remark C.1 (Reduction Single Path GSBM). When (t,o(Xr.o) 0 and (x, 0 for all l E [0, 1] and & E (1,..., K), then the Branched CondSOC problem is the solution to the single path GSB problem given by
[ R P.poo.
( vi(X.) di s.t.
{ [dx, u, (X,))dt odB,
min
Xx ~ Tu. X1 ~ T1
(59)
Mey JO
where the probability density Paj,11 1 (x x1) is conditioned explicitly pair of endpoints (No, x1) ~ T0,1 drawn from the joint distribution.
C.3 PROOF OF PROPOSITION 4.1
Proposition 4.1 (Solving the GSB Problem with Stage 1 and 2 Training). Stage and Stage training yield the optimal drift u. (Xx) that generates the optimal marginal probability distribution Pž (X+) that solves the GSB problem in (3).
Proof. Let the marginal probability distribution pł (X,) and corresponding drift ug (X₂) define the optimal solution the GSB problem. It suffices show that
1. Given (æo, 21) To.1, Stage I training with the trajectory loss (ni(() (12) yields the inter- polant ii,, and time-derivative ii,,, that define pī (X₂).
2. Given P (X²), Stage training with the explicit flow matching loss Lnow (13) yields the optimal drift uz (X.).
To prove Part 1, establish the following Lemma.
Lemma C.1. Given the Markovian reference process Q with drift of (X,) from the unconstrained GSB problem in (3), define the reciprocal projection 11* projR(a) {IP: Po,1 To,1 of the endpoint coupling No,1 onlo the reciprocal class R(Q).. Stage training then learns the velocity field ii.,, whose induced interpolant matches the intermediate time dynamics of from endpoint samples (Io, x1) ~ T0,1-
Proof of Lemma. Since Q C([O, 1]; Rd) is a Markov measure with well-defined bridges 1o. (o = Do, X1 œ1), we define the reciprocal projection II with the endpoint coupling T0,1 as:
rT' -T= f 1.. = Fo X1 = x To, (dwy, dx1)
To.1
(60)
which is in the reciprocal class II R(Q). satisfies the endpoint marginals I7 = To.1, and the unique minimizer of the KL. divergence:
IT
arg min
KL((PIQ)
(61)
PER
Therefore, it suffices prove that ai,,, approximates vf (X.) under endpoint constraints Xo 20 and X, 2 for (1) #0,1-
First, we recall that the unconstrained Markovian drift uy is the minimizer of the energy function given by:
of = arg min / E.. b ve (xt)) v.(x,) dt
(62)
We
where Pt is the marginal law induced by Furthermore, the class of interpolants is given by parameters r is defined as:
Lt.y = (1 -)) + tæi + t(1 - t) (Io. x1)
(63)
i D1 Zo t(1 t) ( (2o, 21) (1 2t) (20, 21)
(64)
27

---

Published as a conference paper at ICLR 2026
Stage 1 training yields the optimal interpolant ii., that minimizes the energy function across all time points t [0, 1I. defined as
[ Epp
1 1 f1 ( x1) + (Z dt
T = arg min
(65)
Therefore, t;,,, matches the reference drift iy defined in (62) along interpolated states while pre-
serving (zo. 21) ~ #0,1- This an upproximation of the bridge dynamics induced by the reciprocal projection IT* (D) (P 1o,1 = _o,1). }. which concludes the proof of Lemma C.
By Lemma C.I, we know that vt is the optimal drift energy function in the GSB problem without end- point constraints and æ2,,, is the approximation of the bridge dynamics induced by II" which follows of while preserving the coupling To Therefore, we can define (xxe/20, æ1) = Ht,, (Zo, æ1) as the conditional drift that generates the conditional probability distribution Pijo. (X, xo. æ1) that satisfies the Pokker-Planck equation
a Ot Proo11 = - (Ptoo) + 2 Aptj,1
(66)
Given that Ii (x, æ1) is the optimal bridge that solves the GSB problem for pair of endpoints (ro; m1) To,1 can define the marginal probability distribution př the mixture of bridges
p = (xo a1) ~ [pio. (X, x1)
(67)
which concludes the proof of Part 1.
For Part 2 of the proof, aim show that Stage training yields the optimal Markovian drift un (X₂) that generates př (X.). To do this, we write the Fokker- Planck equation for pł (X,) in terms of the conditional bridge Prjo,1 and drift field uxjo,1 extract expression for uz (X₂) that satisfies it. Starting from the definition of pł, we have
/ Pijo. dTo.1
p =
0 δe
/
te Pio. doo,1
/ (wHo. + 20
1 o² Apio..1
d
/ (u.o..P(io.)) d + /
1
d
/ (uijo.PPio.. IdTo, +
1
V
022 A
Pijo,dTOo
2
(una.PPioo ) + 702A
pi
(68)
(uTPt)
For (68) to satisfy the Tokker-Planch equation, we set the first term equal to (u(pt) ) and solve for ut
to get
u.p = / d
(69)
Exo. (1
u
(70)
pf
Therefore, the optimal Markovian drift (or Markovian projection) the average of the conditional drifts defined in part as "oo..1 = क over the joint distribution To,1. This means that the minimizer of the conditional flow matching loss Laow (0) in (13) defined as the expected mean-squared error
28

---

Published as a conference paper at ICLR 2026
between a Markovian drift field uy (X²) and ži,, over pairs (Io, x1) No,1 in the dataset is the optimal drift u7 (Xx) that solves the GSB problem.
/ (e ll (Xxell
u (Xe) = arg min (0) = arg min
(71)
which concludes the proof of Proposition 4.1.
Since the drift for each branch u²,x (Xx,x) ))) are trained independently in Stage 2, we can extend this result across all K + 1 branches and conclude that the sequential Stage 1 and Stage 2 training procedures yields the optimal set of drifts {ut,& }K k= that generate the optimal probability paths & that solves the GSB problem for each branch.
C.4 PROOF OF PROPOSITION 4.2
Lemma C.2. Suppose the optimal drift field ui,x R RRd and probability density Pł, Id + R that minimizes the GSB problem in (3) is well-defined over the state space X' C R2 for each branch. Then, the optimal weight wi, any of the secondary branches k E {1.-.., K} is non-decreasing over the interval E [0, 1]. Equivalently, the optimal growth rates gi, (Xt,,) for all [0, 1].
Proef. We will prove this lemma by contradiction. Suppose there exists branch that decreases in weight over the time interval [41,4, 12] for 0 < < l2 S 1, such that wt,, k > Wtg.k- We know that the target weight at time = is non-negative Wi,k 2 and the total weight across all branches is conserved (i.e. E Wr,k = (ntaa) and non-decreasing (i.e. Oeuiotal > 0 for all / E [0, 1)). In this proof, we let & (11, |/" Iu (x),))) I v.((x,)) denote the energy of following the dynamics of the kth branch over the time interval (2).
Then, there can only be two possible reasons for the loss of mass along branch k: (1) the mass is destroyed, or (2) the mass is transferred to different branch j k.
Case The destruction of mass directly violates the assumption that the total mass across all branches is conserved and non-decreasing. So, we only need to consider the possibility of Case 2.
different branch j l over the interval [4, (2). By Proposition 4. Stage 1 and 2 training yields the optimal velocity fields (Xe, that generate the optimal interpolating probability density { (X that independently minimize the GSB problem in (3).
Case 2. Suppose the mass transferred
If mass is transferred from branch i branch j the interval [/,12], 12), it must be compensated for over time [42, 1] to reach the target weight Wi.k > O. Then, without loss of generality, we can consider two sub-cases: (2.1) the mass is compensated from the primary branch, and (2.2) the mass that diverges to the jth branch returns to the kth branch following continuous trajectory.
Case 2.1. Since all mass along the K secondary branches originates from the primary branch, this implies that there exists positive weight Ø > that first follows the dynamics of branch & and transferred to branch j. contributing to the final weight the jth endpoint, and the total weight supplied from the primary branch to branch & is W1,k w. Given that each branch has no capacity constraints, follows that the dynamics along branch & over [0, 4) and branch j over [4, 1] are optimal for all mass reaching endpoint j. Formally, we express this in terms of energy as
E (0, 4) 1), 10,
(72)
which contradicts the assumption that the dynamics of branch j given by (u..))PP.)) optimal over the state space X.
Case 2.2. In this case, there exists a positive mass that follows the dynamics of branch i over the interval [0, ty], the dynamics of branch j over [t, tz], and back to branch & over [t,, 1]. Similarly to Case 2.1. given that each branch has no capacity constraints, this implies that this concatenation of
29

---

Published as a conference paper at ICLR 2026
dynamics is optimal for all mass reaching endpoint k. Expressing in terms of energy, have
(0, f1) + E,(ti, (2) + (t,, 1) < (0,1)
(73)
which contradicts the assumption that the dynamics of branch & given by (ui,.Pi, are optimal over the state space X'
Given that both sub-cases lead to contradiction of the optimality assumption, we conclude that mass along each branch cannot be transferred to another branch and is non-decreasing over t [0, 1].
Note that we do not need to consider the case where the mass compensated from another secondary branch & A j, this would imply that mass is transferred from branch é to branch k, which is not possible under the same argument and we conclude our proof of Proposition 4.2.
Proposition 4.2 (Existence of Optimal Growth Functions). Assume the state space X C R is a bounded domain within I2, Let the optimal probability density of branch k be a known non- negative function bounded in j0, 11, denoted as Pi,x & x [0, i] + [o, 1] E L x [0, 1)).
By Lemma C.2, we can define the set feasible growth functions in the set of gguare-integrabe functions 1.2
= 1g (gr.o..... .t.......999k)) E L 22 [0, 1|: Vkk 2 1, (. 2 Ex= 9 = 03 (18)
Let the growth loss be the functional L(g): L2 x [0, 1K: - Then, there exists optimal function g" 19t,0.-... gt L2 x 10 where 9 E 9 such that
((g") ges L(g) which be obtained by minimizing L(g) over G.
This proof draws several concepts from functional and convex analysis. For a comprehensive background on these concepts, see Benešová & Kružík (2016). We follow the direct method in the calculus of variations (Dacorogna, 1989) for proving there exists a minimizer for the functional f.gwwth (g) in (17) with the following steps:
Proof.
1. Show that the set feasible growth functions is convex and weakly closed under the set of square integrable functions (S2; T2 RK+1) (Lemma C.3).
2. Show that the minimizing sequence (g(*) has weakly convergent subsequence (Lemma C.4).
3. Show that the functional Egowwh(g) is weakly lower semi-continuous (Lemmaa C.5, C.6, C.7, C.8).
We prove each with sequence of Lemmas.
Lemma C.3. The set of feasible growth functions G = {g (gt.o..... E I
[0. 1]: R]: vk 2 1, (x) 2 0 and Ek=o It.k = 0} is convex and weakly closed under
L2 x [0, (++))
Proof. We first prove convexity and then weak closure under 2 10,1);1K+1).
Proof of Convexity: To prove that the set of functions G is convex, we first define what it means for of functions be convex.
Definition C.1 (Convex Set). set of functions G is said be convex if any convex combination of two functions in the set g1, 92 E G is also in the set. Formally, G if
Vg.g' E G, VA e [0, g² = >g (1-A)g e
(74)
Recall our definition of G the set of K + growth rate functions that sum to zero:
G = [g = (g..v....gr.dd
(2 x 10. R [ vk
1, gr, (æ) 2 0 and Σ 9e,4 0} (75)
If g. gf e G and A E [0. 1]. then for k 2 1, we have
Agi (1 2)g' 2 2.0 -) 0
(76)
30

---

Published as a conference paper at ICLR 2026
and adding A O, we have:
Ag + (1 - A)g' EG= (Agio + (1 \)g(,a) + 1
(Agi, + (1
N)
> (9 g0.0 + Σ g6.k) +(1 A) (0i.0 + E + Sk.tk
(77)
which means Ag + (1 - Aig' E G and G is convex.
Proof of Closure. First, we define what it means to be weakly closed under L
Definition C.2 (Weak Closure). A set of functions 9 is said to be closed in the weak topology of L? if the statement is true: if sequence of functions {g : gg()) E } indexed by n, converges to some function g gloo) E T,2 as n - oo, then gloo) Eg.
To show that G is weakly closed under L L2( L2 22 where Ω == X [0, I]. we need to show that all sequences { 9i,k E { converge to g E g as - oc, such that 2
We decompose G := C CnS. where C' is the non-negane non-negative convex cone on the secondary branch growth rates and S is the linear mass conservation constraint, and show that C and S are weakly closed (S2; R K+)). First, we define as:
:= :g (2;; gt.k 2 aac. on X' [kk (L....,,k}
(78)
To prove C' is weakly closed, we need to show that if ((") E C - = C. Suppose g(v) converges weakly to i.e., g(x) - g. Then, for any non-negative square integrable function & E (Q2), clearly the inner product with gim) 2 0 is non-negative:
(g(x)
0
(g.v) 00
(79)
by weak convergence (g("), v) (g. t) Since this holds for all non-negative L2) have that gt.k 2 almost everywhere and g E C, so we conclude that C is weakly closed under (22; +1).
For the linear constraint set S defined as:
S {g L2 (0 R-
90 + Σ 9k = 0 a.c. on X' x [0, 1}}
(80)
Let T: (2; - (22) be the bounded linear operator defined as:
T(g) 9u Ax=1 9k
(81)
Then, have that S (ernel(T)) {g e L ((( T(g) 0} is the kernel of T. To prove weak closure of on ( RK+1), need show for all (x) + - T(g) 0. Since T a bounded linear functional, we have:
g
(( - T(g) (2(
(82)
But we know ((g(n))= 0, which means that:
T(g +0 T(g)= -0
(83)
which proves that S is weakly closed under L L 10; Since we have shown that both C' and S are weakly closed under (2;; IR K+1), we can conclude that their intersection G C weakly closed. This concludes our proof of Lemma C.3.
Lemma C.4. Given a minimizing sequence (g(n) E G} under the functional C(g) L? + R such that (mg(m)) + ing C(g). there exists a subsequence {g(ms) 9G} that converges weakly L to some limit g" G.
Proof. It suffices to show the following:
31

---

Published as a conference paper at ICLR 2026
1. The functional {g(") is bounded in L such that there exists positive value M where [|g(")|1>> I M for all n.
2. Since (S2; IR*+1) is reflexive, every bounded sequence has weakly convergent subsequence.
The proof of Part follows from the growth penalty term in L(g) defined the squared-norm | of the growth functional. This term ensures that L(g) is coercive, such that
(g I|= 00 = ((g!") oc
(84)
which ensures that the sequence does not diverge to infinity in norm without incurring a penalty from the loss functional. Given that [(g(") - inf geg C(g), where gnf L(g) < co, follows from coercivity that (0)(")]|2 I|ا> does not diverge to infinity and [g(0) bounded 1.2.
The proof of Part 2 follows from the well-established result that for < p < co, the space TP is reflexive (Beauzamy, 2011). Therefore, by reflexivity of L2, we have that every minimizing sequence (g(") has a weakly convergent subsequence {g(m))}, such that g(m=) + g" in L concluding our proof of Lemma C.4.
Before proving that each component of the loss functional is weakly lower semi-continuous, we establish the definitions for weakly continuous and weakly lower semi-continuous functionals.
Definition C.3 (Weakly Continuous Functionals). A functional L: L - R is said to weakly continuous it satisfies
n n 7 (n) + - = ((g(n)) f(g)
(85)
such that if a sequence {g(n) g(n) [} converges ((1) + 9 - co, then the functional also converges ((g(n)) + L(g) and g(n)
Definition C.4 (Weakly Lower Cani-Continuous Functionals), A functional C L2 - E is said to be weakly lower semi-continuous (w.lss.c.) if satisfies
n = g( 7 = lim inf (ng ) 2 C(g) n-+0c
(86)
such that if a sequence (n (())) EG} converges ((n)= - as n - oc, then the functional lower bounded by L(g). By definition, weak continuity implies w.l.s.c.
Lemma C.5. The functional Emarch (g) L (2(: - R. defined
Wo.k S: go,dd - -wix 2
Emouc Σ Brice
k=0
wo, + L'''
wix))
(linearity expectation)
k=
Σ ( J pinsnndeed + c)
(87)
where c = Wook wi,k is a constant, is weakly continuous in RK
Proof of Lemma. Define the map:
( ('' pi.
(88)
First, we show that ((gc,) bounded linear functional in (22))). By linearity integration, have that for two functions 9.9 G. $leg + d'g') cΦ(g) d'ó(g'), so Φ(-) is linear
32

---

Published as a conference paper at ICLR 2026
map. To establish that Ø(.) is bounded in (S2), must show that [d(g)] v Cl By the Cauchy-Schwarz Inequality (Steele, 2004), we have
[ ( 0t,aa1 L.² [. pi. Ippi. |epi:lle
(89)
which is valid given pi At.k (9)). Since I ((a)) is fixed constant with respect 9t,k, have shown that [d(g(x)) Cl and It.k is bounded. By the definition of weak topology on (22), all bounded linear functionals are weakly continuous, such that limm-so ((g(") + d(g).
Given that ((g) is weakly continuous, as g(m) - gg we have ((g()) + "(g"). ). Since the square function (-()) (-)² is convex, continuous, and bounded below by v O, the function does not contain discontinuities and we have weak convergence of (Ø(g(")) c) - - Therefore, we can conclude that each term the sum weakly continuous, and since sum of weakly continuous functionals is also weakly continuous, conclude proof of Lemma C.5.
Lemma C.6. The functional Lenergy (g) 1.2 RS;; - defined as
Cmergy (g) = S/' Epix (2l +i (wo.++ + J(& (h,e(X
vi(Xee)) = S.' a(t) wo.k + So" PasDerddadss dt
(lincarity of expectation)
where att) is constant not dependent on It,k, is weakly lower semi-continuous in L (22 R 11) (w.ls.c.)..
Proof of Lemma. For each E [0, 1]. we define the functional:
(.(gx,)) / Po.
(90)
which is a bounded linear functional on (2) for all s. Applying Fubini's theorem to change the order of integration, we have:
o')(t) a(t) (fee.ooas)) dt ( a(t)att o
(91)
=;((x)
Defining B(6) = S. a(l)dl, we have:
S. ((s) / Parg (8p% gb) L2(0))
(92)
which is bounded linear functional in I and thus is weakly continuous in L and we have shown that Lenergy (g) is weakly continuous. Since weak continuity implies w.l.s.c., and the sum of w.l.s.c. functionals also w.l.s.c., we conclude our proof of Lemma C.6.
Lemma C.7. The mass loss functional Emase (g) (2( - R defined as
2
Limrr (g) = [f. BR M
Σ W
notal
dt
K
Σ
Wu.k +
1/' Eps, (ocxal)
w
(lincarity expectation)
kmm
K
P
nofttal dt.
(93)
K=0
where wootal is constant not dependent It.k. is weakly lower semi-continuous in [2( (w.lssc..). Note that do not need to account for the negative penalty loss as we assume the sum of growth functions is zero.
33

---

Published as a conference paper at ICLR 2026
Proof of Lemma. For this proof, consider the bounded linear operator A (2( (2 1) defined as:
K (Ag)(t) = Σ
J( Pa.N (.. K=U
(94)
Then, we have:
C (g) ! (Ag)(t) - w/ 2 dt |1 (Ag)
wi [otaall (
(95)
Since the bounded linear operator A is weak-to-weal continuous and the squared norm || [2 is w.l.s.c. on 1.2 (10, 1)). g(w) in (22; 11), then Ag Ag [[, 1]) and we have:
lim Lmaoo (g (m)) = lim inf n=000
+ Ag (m)
wi (.P(() 2 1 Ag wi (2((0,1)) Cmis (96)
n-
which means Lmass (g) is w.l.s.c. and we conclude our proof of Lemma C.7.
Lemma C.8. The combined Stage 3 training loss Egrownt (g) Xenergy energy (g) +
Amatch mmarch (g) + Amass ( + Xxrowwhh Illll; is weak lower semi- cemi-continuous in ( R K+ +1) (w.ls.c.).
Proof of Lemma. Given the weighted sum of w.l.... functionals Σ A,Co(g) that each satisfy oim Li(g()) Li(g) for constants {M follows easily that
JM for converging sequence g(²) 9 oc,
M Σ (.( 2 Σ xg)(g)
M
lim inf +20
(97)
1=1
where the weightes A; 2 O. By definition, the L norm is weakly lower semi-continuous. Since Cunergy fmatch- and Cuuass are w.l.s.c. by Lemmas C.5, C.6, and C.7, we conclude that Cgrowth (g) is w.l.s.c., which concludes the proof of Lemma C.8.
In total, we have shown:
1. The set G of feasible growth functions is convex and closed under the weak topology of the set of square integrable functions 2 (Lemma C.3).
2. The minimizing sequence (g(")} has weakly convergent subsequence (Lemma C.4).
3. The functional L(g) weakly lower semi-continuous (Lemmas C.5, C.6, C.7, C.8).
Thus, by the direct method in the calculus of variations, we have shown
3g = (gi gi,1) E G s.t. L(g") inf ((g)
(98)
gEG
which concludes the proof of Proposition 4.2.
D ADDITIONAL EXPERIMENTS AND DISCUSSIONS
D.1 COMPARISON TO BRAGNH SCHRÖDINGER BRIDGE MATCHING
Setup For each experiment, we compared the performance of BranchSBM against single-branch SBM. Instead of clustering the data data as a single target distribution the initial distribution No- For the single-brand branch task, we assume mass conservation and set the weights for all samples from To and n1 to 1.0, while keeping the model architecture, state-cost Vz, and hyperparameters equivalent to BranchSBM. Since single-branch SBM does not require modeling the growth of separate branches, we train only Stages I and 2 to optimize the drift field uf of the single branch. For evaluation, the trajectories of validation samples from the initial distribution Tp To were simulated over t [0. 1] and compared with the ground truth distribution at 1. For BranchSBM, we take the overall distribution generated across all branches and compare to the overall ground truth distribution.
into distinct endpoint distributions, we left the unclustered and let the model learn the optimal Schrödinger Bridge from
34

---

Published as a conference paper at ICLR 2026
Principal Components
100 Principul Components
150 Princjus Componentt
Proepd Cr900e13
Single Branch SBM PCS
Figure 6: Results for Clonidine Perturbation Modeling, (A) Gene expression data of DMSO control (set to t == 0) and cell states (set to 1) after Clonidine perturbation with two distinct endpoints (pink and purple). (B) The simulated trajectories for single-branch SBM 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 & = on the validation data for each branch.
Modeling Mouse Hematopoiesis Differentiation In Figure 7, we provide side-by-side compari- of the reconstructed distributions at time L and l2 using BranchSBM (top) with branches and single-branch SBM (bottom), as well as the learned trajectories over the time interval / E [4, 12]. While single-branch SBM produces samples that loosely capture the two target cell fates, the resulting distributions display high variance and fail to align with the true differentiation trajectories. In contrast, BranchSBM generates intermediate distributions that are sharply concentrated along the correct developmental paths, more faithfully reflecting the underlying branching structure of the data.
A
Simulated State at Time 11
B
Simulated State Time 12
Learned Trajectories
BranchSBM
Single Branch SBM
PC1
Figure Comparison BranchSBM to Single-Branch SBM for Cell-Faae Differentiation. Mouse hematopoiesis scRNA-seq data is provided for three time points fo, f1, f2. (A, B) Distribution of simulated cell states at time (A) and (B) across both branches for BranchSBM (top) and single-branch SBM (bottom). (C) Learned trajectories of BranchSBM and single-branch SBM over the i [4, te] validation samples.
Modeling Clonidine and Trametinib Perturbation For both Clonidine and Trametinib, we performed the single-branch experiment the top 50 principal components identified by PCA. After training, we simulated all the validation samples from the initial distribution To by integrating the single velocity field of to l 1. We evaluated the performance of the single-branch parameterization
35

---

Published as a conference paper at ICLR 2026
by computing the RBF-MMD (104) of all PCs and W1 (102) and W2 (103) distances of the top-2 PCs of the simulated samples at time = 1 with the ground truth data points.
In Table 4, we show that BranchSBM with two branches trained on gene expression vectors across all dimensions d E (50, 100, 150) outperforms single-brane SBM on dimension d = 50 recon- structing the distribution of cells perturbed with Clonidine. In Table 5, further show improved performance of BranchSBM with three branches compared to single- SBM.
In Figure 8A and 8B, to generate samples from cluster for Clonidine perturbation or clusters and 2 for Trametinib perturbation. The endpoints for both perturbation experiments first two or three principal components (PCs), where PC1 captures the divergence from the control state the perturbed clusters and PC2 and higher captures the divergence between clusters in the perturbed population, where cluster 0 is closest to the control state along PC2 and cluster 1 and are farther from the control. From Figure 8A and B, we can conclude that single-branch SBM not expressive enough capture the complexities of higher-dimensional PCs and follows the most obvious trajectory from the control to cluster O, resulting in an inaccurate representation the perturbed cell population.
that single- branch SBM only reconstructs cluster 0, while failing
clustered largely variance in the
In contrast, we demonstrate that BranchSBM is capable stimulating trajectories to both clusters in the population perturbed with Clonidine (Figure 8B) and three clusters in the population perturbed with Trametinib (Figure 8D), generating branched distributions that accurately capture the location and spread of the perturbed cell distribution in the dataset.
Clonidine Perturbation (2 Endpoints)
Trametinib Perturbation (3 Endpoints)
hapcter start Txectory Ord
Tepectary Start Tajjctory End
Single Branch SBM
BranchSBM
Figure Comparison BranchSBM to Single-Branch SBM for Perturbation Modelling. (A, B) Clonidine perturbation trajectories with two target clusters generated by (A) single- branch SBM and (B) BranchSBM from the validation data. (C, D) Trametinib perturbation trajectories with three target clusters generated by (C) single-branch SBM and (D) BranchSBM.. In both experiments, single-brancl SBM only generated states cluster and cluster or 2, whereas BranchSBM reconstructed all perturbed clusters via branched trajectorics.
36

---

Published as a conference paper at ICLR 2026
Pancreatic B-Cell Differentiation
Branch?
Tiee Componer
Branch
Branch 10
Branch 11
All Timepoints
Figure Results for Panercatic 3-Cell Differentiation Modeling with BranchSBM. We evaluate BranchSBM on pancreatic 3-cell differentiation dataset containing 51, 274 cells collected across eight time points. Cells are projected into 30-din do-dimensional PCA space, and Leiden clustering is used to define 11 terminal cell populations the final time point. The trajectories of each branch & E {0, .... 10} simulated from validation points and are plotted with the top-2 PCs. The bottom-right plot shows the cells from each intermediate time point.
D.2 EFFECT OF FINAL JOINT TRAINING ON LOSSES
In Table 7, show the final losses after convergence, summed across each branch and averaged over the batch size, for Stage 3 training of only the growth networks and Stage 4 joint training of the flow and growth networks discussed Section 4. All losses are calculated exactly shown in Section 4. Crucially. we find that the final joint training stage refines the parameters of both the flow and growth networks simultaneously to minimize the energy loss Cenergy (9, o) (14) while ensuring the growth parameters maintain minimal losses across Lematch (Φ) (15) and Lmmss (16) for all experiments. This indicates that jointly optimizing both the drift and growth dynamics leads to further refinement towards modeling the optimal branching trajectories in the data.
Table Validation Losses Stage and Training Aeross Experiments. Losses are summed across both branches and averaged over batch size. The final Stage 4 joint training stage that refines both the flow and growth networks sinultaneously minimize the energy loss Lesengy (0, Φ) (14) from Stage 3 while ensuring the growth parameters maintain minimal losses Lmauch (6) (15) and Lamss (16)
Experiment
Stage 3
Stage 4
Levery (0,
Cmm ())
Lmmeh Φ))
Levery (0, ())
Cmm.0 (0,
Emareh (9,
LiDAR
1.276
3.0 x 10-5
0.007
0.768
2.0 10-5
0.102
Mouse Hematopoiesis
2.209
1.2 x 10-4
0.054
1.918
5.0 10
0.076
Chlonidine Perturbation
36.469
0.030
0.109
25.798
0.053
0.153
Trametinib Perturbation
35.834
0.023
0.078
32.843
0.017
0.056
37

---

Published as a conference paper at ICLR 2026
E EXPERIMENT DETAILS
E.1 MULTI-STAGE TRAINING
To ensure stable training while incorporating all loss functions, we introduce a multi-stage training approach (Algorithm 1).
Stage 1 First, we train a neural interpolant (P,,, (oo, Ft,s)) RR Rd x [0, - g that takes the endpoints of branched coupling and defines the optimal interpolating state X, at time by minimizing the energy function (nn in (12). This used to calculate the optimal conditional velocity &r,n, that preserves the endpoints Xo 2g and X1,6 Ii for the flow matching objective in Stage 2.
Stage 2 Next, train set of flow networks {ul, k=0 that generate the optimal interpolating trajectories for each branch with the conditional flow matching objective Lnow in (13).
Stage 3
We freeze the parameters of the flow networks and only train the growth networks (9 k= by minimizing Lgowth in (17).
Stage 4 Finally, we unfreeze the parameters of both the flow and growth networks and jointly train fu Fi by minimizing the growth loss Lguwith (17) from Stage 3 in addition the distribution reconstruction loss Creonss in (19).
E.2 GENERAL TRAINING DETAILS
Leiden Clustering To identify branch endpoints in the dataset, we apply an automated clustering pipeline to the final time point (t = 7), where the differentiated cell-states are most clearly defined. First, construct
neiichb (kNN) graph the final- imal-timepoint embeddings and run Leiden community detection using the Modularite wweettteeeaaareeee tytee clustering is particularly suited and the state-of-the- method for clustering single-cell data, given that it is robust to noise, can handle heterogeneous cluster sizes and shapes, and guarantees well- connected clusters, all while being fully unsupervised and automatic. This gives us stable and biologically meaningful terminal clusters compared to alternative methods such as K-means clustering, which assumes spherical clusters and struggles with sparse -igh-iimensi manifolds. After initial clustering, we merge any cluster with fewer than _ellss cells into the nearest large cluster based on Fuclidean centroid distance.
on objective. Leiden
Model Architecture We parameterized the branched trajectory (Pt, (2o. æ1) with a 3-layer MLP with Scaled Exponential Linear Unit (SELU) activations. The endpoint pair (æo, x1) and the time step are concatenated into single (2d 1 1)-dimensional vector and used as input to the model. Similarly, we parameterize each flow network uº (xx) and growth network gf, (xx) with the same 3-layer MLP and SELU activations but takes the interpolating state x, and time concatenated into (d + -idimensio vector.
To ensure that the growth rates of all secondary branches are non-negative (i.e. for all k: E (1...., K}. It,x (X,,) 2 0), we apply an additional softplus activation to the output of the 3-layer MI.P in the growth networks, defined softplus(-) log(1 1 exp(-)), which is smooth function that transforms negative values to be positive near O. For the growth network of the primary branch (k = 0), we allow for both positive and negative growth, as all mass starts the primary branch and flows into the secondary branches, but the primary branch itself can grow as well, depending on whether mass conserved.
State Cost Vj Depending on the dimensionality of the data type, we set the state cost Vi(X,) Rd + [0. +oo| to be either the LAND RBF metric discussed in Appendix A.2. For the experiments on LiDAR (d = 3) and Mouse Hematopoiesis scRNA-seg (d = 2) data, we used the LAND metric (33) with hyperparameters
0.125 and 0.001.
38

---

Published as a conference paper at ICLR 2026
To avoid the task of setting a suitable variance for the high-dimensional gene expression space E {50, 100, 150). we use the RBF metric (37) that learns parameters to ensure the regions within the data manifold have low cost and regions far from the data manifold have high cost. Using the training scheme in Kapušniak et al. (2024), we identified N₂ cluster centers with k-means clustering, and trained the parameters {wj.n N. by minimizing LRBE (38) on the training data. We found that setting the number of cluster centers No too low resulted in non-decreasing CRBF and that increasing Ne for higher dimensions enabled more effective training. Furthermore, we found that modeling sigher-dimensionll PCs required setting a larger K, which determines the spread of the RBF kernel around each cluster center. The specific values for N. and depending on the dimension of principal components are provided in Table 11.
Optimal Transport Coupling Since our experiments consist of unpaired initial and target distribu- tions and we seek to minimize the energy of the interpolating bridge, we define pairings (æo, x1.k) using the optimal transport plan 10.1 that minimizes the distance between the initial distribu- tion To and each target distribution & in probability space, Specifically, we define T6,1. as the 2- Wasserstein transport plan (Villani et al., 2008) between To and T defined as
/
L = arg min min L..L
I00 21, x1.x (wo, 21,kk
(99)
where To & T is the set of all possible couplings between the endpoint distributions. For each of the branches, the dataset was paired such that (æo, x1,,)) ~ T..
Training We train each stage for maximum of 100 epochs. For Stage 1, we used the Adam optimizer (Kingma& & Ba, 2014) with a learning rate of 1.0 x 10 to train (2,, æ1). For Stage 2, 3, and 4, we used the Adam W optimizer (Loshchilov & Hutter, 2017) with weight decay 1.0 x 10-5 and learning rate 1.0 x 10-3 to train each flow network uf, and growth network gi. All experiments were performed on one NVIDIA A100 GPU. We trained on the LiDAR and mouse hematopoiesis data with batch size of 128 and the Clonidine and Trametinib perturbation data with a batch size of 128, each divided with 0.9/0.1 train/validation split. All hyperparameters across experiments are provided in Table 11.
Computational Overhead Although we train K 1 flow and growth networks, the overall training time remains comparable to that of annle-branch SBM, since the networks for each branch are trained only on the subset of data corresponding to its respective target distribution. While the method does incur higher space complexity, we find that simple MIP architectures suffice for strong performance, suggesting that scalability is not a major concern. BranchSBM also significantly reduces inference time, as predicting branching population dynamics requires simulating only single sample from the initial distribution, unlike other models that require simulating large batches samples.
E.3 LIDAR EXPERIMENT DETAILS
LiDAR Data We used the same LiDAR manifold from Liu et al. (2023a); Kapušniak et al. (2024). The data is collection of three-dimensional point clouds within 10 unit cubes [-5,5] C R that span the surface of Mount Rainier. Given any point æ E R3 in the three-dimensional space, we project it onto the LiDAR manifold by identifying the k-nearest neighbors (.1,..., and fitting a 2D tangent plane to the set of neighbors
arg min b.
Σ exp K
1
(
xill| ( -xill
(aar(")
+ banl") + - {2
(100)
where - = 0.001 following Liu et al. (2023a). Then, we solve for the tangent plane ar + by + c = using the Peorree pseudoinverse from k = 20 neighbors. From the tangent plane, we can project any point æ to the LiDAR manifold with the function <ι)) defined as
x v +c)
7 (Φx)
v, where
b
(101)
||v
-1
39

---

Published as a conference paper at ICLR 2026
Synthetic Distributions To reformultte the experiment in Liu et al. (2023a) as branching problem, we define a single initial distribution and two divergent target distributions. Specifically, we define a single initial distribution No N (( σo) mixture of four Gaussians and two target distributions 11,0, T1,1 on either side of the mountain, both as mixtures of three Gaussians. The exact parameters of each Gaussian are given in Table 8.
We sample a total of 5000 points i.i.d. from each of the Gaussian mixtures [ {0i 15000 11,0, and {=i_ 15000 TI,1. The data points are projected to the LiDAR mani- fold with the projection function (ax) in (101).
1,500
To
Table 8: Synthetic Gaussian mixture distribution parameters for LiDAR experiment. 5000 datapoints are drawn i.ddd from each of the Gaussian mixtures and paired randomly (wo. 21,1) to define the training datasct. A visualization of the training data the LiDAR manifold provided in Figure
Distribution
d
50
M(uo .o) (1 F1.))
(-1.5, -..0 (-1.2, 3.5, (-4.0, 3.0, (-3.75, -2.5, 0.5)
0.02
$1.0
(-2.25, 0.675, 0.5). (-2, 15,0.) -0.25, (-25,-0.2,,,,),
0.03
(2, -2,0.), (2.6, -1.55, 0.5), (3.2, 0..5
0.03
To determine how closely the simulated trajectories match the ground truth trajectories, we compute the Wasserstein-1 (W) and Wasserstein-2 (W2) distances defined as
Evaluation Metrics
/ 11-ylldd((aa
Wi
min xEIl(p.9)
(102)
/ 1x-yyll y (z.
1/2
W2
min (xll(pp.)
(103)
where p denotes the ground truth distribution and y denotes the predicted distribution. After training the velocity and growth networks on samples from the initial Gaussian mixture To and target Gaussian mixtures 11,0 and 11.1. evaluate Wi and W/ of the reconstructed distribution simulated from the validation points in the initial distribution Tg against the true distribution at =
E.4 DIFFERENTIATING SINGLE-CELL EXPERIMENT DETAILS
Mouse Hematopoiesis scRNA-seq Data We used the mouse hematopoiesis dataset from Zhang et al. (2025c) consisting three timesteps to: f1, 12, with a total of 1429 cells at tn. 3781 cells from h, and 5788 cells from 12. The data points at tp form homogeneous cluster, while the data points at 1/ are clearly divided into two distinct cell fates. We performed k-means clustering with clusters create branching on the cells 1g. splitting the cells into clusters: endpoint 0 with 2902 cells and endpoint with 2886 cells. Since the two endpoints are near equal in ratio, we set the final weights of both endpoints as 0.5 (i.e. W1.0 W1.1 0.5). To match the size of the ty samples, we randomly sampled 1429 samples from each of these two clusters and used them as the endpoints for branches 0 and respectively. Training and validation follow a 0.9/0. split ratio.
Pancreatic 3-Cell Differentiation Data The pancreatic B-cell dataset from Veres et al. (2019) contains 51, 274 cells over eight time points that evolve from human pluripotent stem cells to pancreatic B-like cells. Following (Zhang et al., 2025d), the gene expression representations are projected to a 30-dimensional PCA space. After Leiden clustering with A = 20 and min_cells- 100, we obtain 11 clusters for K 11 branches.
Evaluation Metrics To determine how closely the reconstructed distributions along the trajectory match the ground truth distributions, we compute the 1- 1-Wasserstein (Wr) (102) and 2- 2-Wasserstein (W2) (103) distances similar to the LiDAR experiment. After training the velocity and growth networks samples from the initial cell distribution TT, and differentiated target cell distributions Tta.0 and 12,, we evaluate Wi and W2 between the reconstructed branched distributions at the intermediate time t1 (Pr,, o and Pt,. and the target distributions at the final time ty (Pt. and Pt,.1) simulated from the validation points in the initial distribution To and the true distributions T,, and Ttz-
40

---

Published as a conference paper at ICLR 2026
Time (1429 Cells)
Time (3781 Cells)
Time (5788 Cells)
1.3
-13
235 1.5 0.5
-35 1.5 -1.
10
Figure 10: Mouse Hematopoiesis Single-Cell RNA Sequencing Data Plotted Time Point. Real scRNA-seo data is projected to 2D force-directed SPRING plots (Sha al., 2023; Weinreh et al., 2020; Zhang al., 2025c). There a clear divergence
cell fate between times fo (left), & (middle), and f2 (right) from the initial homogeneous progenitor cells into two distinct cell fates (shown in pink and purple in the t2 plot). Cells at time l2 are clustered into endpoint 0 (pink; 2902 cells) and endpoint 1 (turquoise: 2886 cells).
E.5 CELL-STATE PERTURBATION MODELING EXPERIMENT DETAILS
Tahoe Single-Cer Cell Perturbation Data The Tahoe-100M dataset consists of 50 cell lines and over 1000 different drug-dose conditions (Zhang et al., 2025a). For this experiment, we extract the data for single cell line (A-549) under two drug perturbation conditions selected based cell abundance and response diversity.
Clonidine at 5 μL was selected first due having the largest number of cells at this dosage, while Trametinib was chosen as the second drug based on its second-highest cell count under the same con- dition. For both drugs, selected the top 2000 highly variable genes (HVGs) based normalized expression and projected the data into a 50-dimensional PCA space, which captured approximately 38% of the total variance in both cases.
We applied K-nearest neighbor (K = 50) and conducted Leiden clustering separately for drugged versus DMSO control conditions. The most abundant DMSO cluster was selected as the initial state (t = 0). For Clonidine, we identified two clusters that were most distinct from the DMSO control along PC1 and PC2, respectively. These were selected as two distinct endpoints for branches and 2 at f = 1. We applied centroid-based sampling to obtain balanced training sets of 1033 cells per cluster (Figure 11).
For Trametinib, we extended the branching up to three endpoints. From its Leiden clustering results, we identified three clusters that were the most divergent from the DMSO control clusters along PC1, PC2, and PC3, respectively. All selected clusters contained at least 100 cells and were subsampled to 381 cells each for branch training. The remaining cells were clustered with K-means into three (Clonidine) four (Trametinib) groups to construct the metries dataset. The training and validation dataset split followed a 0.9/0.1 1 ratio. The final visualization utilized the first two principal components.
Table 9: Training cluster cell counts for perturbation experiments.
Clonidine
Trametinib
Cluster
Cluster 1
Cluster
Cluster
Cluster 2
Original Cell Count
1675
1033
1622
686
381
Initial Weight Wa.k
1.0
1.0
Target Weight w1,e
0.619
0.381
0.603
0.255
0.142
Evaluation Metrics To quantify the aligmment of the reconstructed and ground-truth distri- butions for the cell-state perturbation experiment principal component (PCs) dimensions
41

---

Published as a conference paper at ICLR 2026
Clonidine Perturbation Dataset
Trametinib Perturbation Dataset
treweet treweet
trawers
triwe11
Compontons
Pincipal Component
Pric pai Component
Figure 11: Clustered Cell-State Perturbation Data from the Tahoe- 100M Dataset. PCA was conducted on all cells for the control DMSO-treated and drug-treated populations, and clustered. Plots show divergence along the top PCs. (A) Clonidine-treated cells (5μM) are plotted pink (endpoint 0), turquoise (endpoint 1), and the distribution control cells is plotted in navy. (B) Trametinib- cells (5μM) are plotted in purple (endpoint 0). turquoise (endpoint 1), and pink (endpoint 2). and the distribution of control cells is plotted navy.
E {50, 100, 150}, we calculate the Maximum Mean Discrepancy with the RBF kernel (RBF-MMD) on all predicted PCs and the 1-Wasserstein (W1 (102) and 2-Wasserstein (W2) (103) distances the top-2 PCs.
This choice is made because the Wasserstein distance between empirical measures (i.e., discrete point clouds) is known to converge to the true Wasserstein distance between the underlying distributions at a rate that scales as & = On -1/d) for some constant C, where n is the number of samples and is the dimensionality of the state space. Therefore, exponential number of samples is required to achieve the same estimation accuracy in high dimensions.
Computing the Wasserstein distance on the top-2 principal components provides a statistically stable and computationally tractable approximation that preserves the major axes of variation in the data. In addition, evaluate the RBF-MMD metric all simulated PCs, ensuring that higher-dimensional reconstruction accuracy is still quantitatively assessed.
Given the predicted distribution p and true distribution g and n samples from each distribution {at p and {y: D the RBF-MMD between p and q is calculated as
MMD(p, 4)
1
Σ
Kiiis (Ie,, 2;)) )
Jmmix (yr,ys)
2 n
Kiiix (w ys)
(104)
1=1
where himix (..) is mixture of RBF kernel funetions defined as
kmix (x, y)
1 M
Σ OEE
exp
[ y 11
(105)
20
where Σ (0.01, 0.1, 10, 100) is the set of values that determine how much the distances between pairs of points are scaled when computing the overall discrepancy. The equations for I-Wasserstein (W1) and 2- Wasserstein (W2) distances are provided in (102) and (103) respectively.
E.6 COMPUTATIONAL COSTS
For each experiment, we report the GPU hours on a single NVIDIA A100 GPU, demonstrating that only minor increase in compute time is required for increasing the number of branches.
42

---

Published as a conference paper at ICLR 2026
Table 10: Total training time for each experiment NVIDIA A100 GPU for different dimensions and number
branches
Experiment (Branches)
Total Training Time (min)
Mouse Hematopoiesis (K
13m 4s
Mouse Hematopoiesis (K
14m 3s
Clonidine S0D (K = 1)
17m 22s
Clonidine S0D (K 2)
22m 50s
Clonidine 100D (K 2)
23m 3s
Clonidine 150D (K" 2)
18m 31s
Trametinib (K = ()
4m $2s
Trametinib (K 3)
6m 44s
Pancreatic 33-Col (K 11)
75m 46s
E.7 HYPERPARAMETER SELECTION AND DISCUSSION
In this section, present the hyperparameters used in each experiment. While the model architecture remained largely the same across experiments, we increased the hidden dimension to 1024 for di- mensions d E {50, 100, 150). For low-dimensional data, found that increasing model complexity underperforms in comparison to lower hidden dimensions, and we established that hidden dimension of achieves relatively optimal performance for d E 3). While beyond the scope of this study, we believe that further exploration of diverse model architectures and hyperparameter tuning could improve the performance BranchSBM. Exploration of diverse task-dependent state costs for novel applications is another exciting extension of our work.
Table 11: Hyperaarameter settings for different datasets. The Clonidine perturbation experiment is split into three columns for cach of the threc dimensions of principal components (PCs) used E {50, 100, 150).
Parameter
Dataset
LiDAR
Mouse Hematopoiesis scRNA
Clouidine 100PCs
Trametinib
Pancreatic fi-Cell
50PC's
150PCss
hranches
data dimension
3
2
50
100
150
50
30
batch size
128
128
32
32
256
1.0
1.0
1.0
1.0
1.0
Araus
100
100
100
100
100
Drzzck
102
10'
10
10'
Anam
1.0
1.0
1.0
1.0
1.0
Agranh
0.01
V.
LAND
LAND
RBF
RBF
RBF
RBF Nr
150
300
300
150
300
RBF
1.5
2.0
3.0
1.5
3.0
hidden dimension
64
64
1024
1024
1024
Ph.n
10
10
10
10
Ir u
1.0 x 10
1.0 10
1.0 10
1.0 10
1.0 10
Ir o7
1.0 x 10-3
10 10
1.0 10
10 10
1.0 10
43

---

Published as a conference paper at ICLR 2026
F TRAINING ALGORITHM
Here, we provide the pseudocode for BranchSBM's multi-stage training algorithm for stable opti- mization of the velocity and growth networks over the K branched trajectories.
Algorithm 1 Multi-Stage Training of BranchSBM
1: Stage 1: Learning the Branched Neural Interpolants
2: while Training
3: Vk, (Zo, #11,k) ~ "0 t U(0,1)
4: for k = 0 to K do
5: Ft.nkk f (1 - 2) + 10 + t(1 P) ( (1,k)
6:
Lt.n,k
+ 1(1 - ( x 21) (wo; æ1,a)
7:
Compute Vi (2 given the task-specific definition
8:
(1)j(n) Jo [ ,, | + V.(( de
9: Update PP., using gradient V v,,Cunj (1)
end for
11: end while
Stage Initial Training of Velocity Networks
13: while Training do
14:
Initialize + flow networks (Xtx)))
for to K do
16:
Calculate Zt,y,k and Lt.nnk with the trained network pi,,, (wo, m1,k)
CRow || ( Update ni,x using gradient V VoLnow end for
18:
19:
20:
end while
21:
Stage 3: Initial Training of Growth Networks
22: while Training do
Freeze parameters of flow networks and initialize K + growth networks ( (x,,,))
for 0 do
25:
for k = 0 to K do
Ft,h +- So (I
if k: 0 then
28:
win + 1 + So 9i,k (x
29:
else
with So (2
end if
Lenergy (0) Lenergy (Ø) + J - 2 All vi(Xe.t) w.
33:
end for
Liass + (IK=e "i, - -ootal) 2
34:
35:
end for
Lmatch f- Ek=0 (Wf (21,)) wT,x
36:
37:
Lrecons (0) +- ER= Ea, LEN max(O, ||āi, - 21, -c)
38:
Leronthh (0) + >energy Fenogy (0, Φ) + Amotch Smarch (0) + Amass Emass (0) + Agrowth ΣR 11 113
39:
Update 9h using gradient V (grown)))))
40: end while
41:
Stage 4: Final Joint Training
42: while Training do
43:
Unfreeze parameters of flow networks (X,,x(Xe,))
44: Repeat steps of Stage 3 and caleulate Ljoin (0, .)) x- Legrowth (0, Φ) Lrooos
45:
Jointly update Uf, and 9h, for all branches using gradients V (0, x) and V V (f,
46: end while
44

---

Published as a conference paper at ICLR 2026
G USE OF LARGE LANGUAGE MODELS (LLMS)
We acknowledge the use of large language models (LLMs) to assist in polishing and editing parts of this manuscript. LLMs were used refine phrasing, improve clarity, and ensure consistency of style across sections. All technical content, experiments, analyses, and conclusions were developed by the authors, with LLM support limited language refinement and editorial improvements.
45


Reviews:
[{"review_id": "JzY3fjgSwI", "reviewer": "Reviewer_3tYX", "summary": "The paper introduces Branched Schrödinger Bridge Matching (BranchSBM) to overcome the limitations of existing approaches, such as flow matching and Schrödinger Bridge Matching, which can only model a single unimodal path. BranchSBM learns multiple time-dependent velocity fields and growth processes, enabling it to capture branched or divergent evolution from a common origin to multiple distinct outcomes.", "questions": "Could this method be applied to tasks such as text-to-image generation or other common multimodality problems? If so, is there a reason why the paper does not explore these applications? \n\nPlease also refer to the weaknesses section for additional context.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "* This method is grounded in solid mathematical theory and is particularly useful for challenging tasks, such as multi-path surface navigation.\n* The paper adopts an approach similar to OT-CFM, where optimal transport is used to define couplings that improve flow-matching vector fields. It formulates the **Branched Generalized Schrödinger Bridge Problem** and develops a theoretical foundation around it, enabling the computation of couplings for downstream CFM models to learn branched paths effectively. This is a common method, yet its deduction is non-trivial in this case.", "weaknesses": "There is room to improve the clarity of the technical sections by providing more detailed explanations. For example, could the authors clarify why it is necessary to consider the time-dependent weights and growth rates in Equation (5)? What is the motivation or scenario that calls for this addition, and how does it connect to the problem studied in the paper? Explaining why this is the appropriate formulation would help readers better understand both the method and the rationale behind it.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oLr5ofXIa9", "reviewer": "Reviewer_7Why", "summary": "This paper introduces Branched Schrödinger Bridge Matching (BranchSBM), a generative framework designed to model transport from a single initial distribution to multiple distinct terminal distributions. Current models, like standard Schrödinger Bridge Matching (SBM) and flow matching, are limited to modeling single, continuous paths between a source and target. The goal of BranchSBM is to learn branched trajectories to guide from a single source to multiple target by parameterizing separate dynamics for each path. The method also extends the Generalized SBM (GSB) framework by including a non-linear state cost, $V_t(X_t)$, which acts as a potential term to ensure that the learned trajectories remain on the data manifold.\n\nThe proposed BranchSBM algorithm solves the \"Branched GSB problem\" by parameterizing separate velocity fields (drifts) $u_{t,k}$ and growth rate fields $g_{t,k}$ for each of the $K$ branches with neural networks. This formulation allows the model to handle \"unbalanced\" transport, where the mass of each branch can grow or shrink. The model is trained using a four-stage algorithm. Stage I, trains a neural interpolant to find the optimal, energy-minimizing paths between the source and each target endpoint. Stage II then uses conditional flow matching to train the branch-specific drift networks $u_{t,k}$ to replicate the velocities from Stage 1. Stage III freezes the drift networks and trains the growth networks $g_{t,k}$ to match the known target mass of each branch. Stage IV unfreezes all parameters and jointly fine-tunes all networks to minimize a combined objective that includes energy, mass matching, and distribution reconstruction losses.", "questions": "The method's parameters and computation scale linearly with the number of branches, $K$. You've shown this for $K=2$ and $K=3$. Have you tested the method's stability and performance for a larger $K$, such as $K=10$ or $K=20$, which would be more representative of complex biological differentiation?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The paper tackles an interesting and well-motivated problem in machine learning of how can we model population-level dynamics that diverge from a single source to multiple distinct outcomes. This is common scenario in fields like single-cell genomics, where we are interested in how a single, source population of pluripotent cells evolved into several distinct cell types, and standard generative models that assume a single, unimodal transition are insufficient. The authors' core idea of extending Schrödinger bridges to explicitly handle branched paths is a relevant contribution.\n\nThe theoretical framework for decomposing this complex problem also appears sound. The authors provide proofs for their core propositions, notably Proposition 1, which tractably reframes the \"Unbalanced GSB\" problem into a solvable \"Unbalanced CondSOC\" objective, and Proposition 2, which then justifies modeling the full \"Branched GSB\" problem as a sum of these individual objectives . This provides a formal grounding for the method's architecture, which separates the dynamics for each branch.", "weaknesses": "The proposed method's significant complexity raises concerns about its practical adoption, scalability, and generalizability. The algorithm is a four-stage sequential training pipeline (Algorithm 1) that involves training at least $2(K+1) + 1$ separate neural networks (an interpolant, $K+1$ drift fields, and $K+1$ growth fields). Furthermore, the overall objective is a carefully weighted sum of at least 6 different loss terms (trajectory, flow, energy, match, mass and reconstruction), which are balanced by many hyperparameters. This multi-stage, multi-loss design suggests a high sensitivity to tuning and raises questions about its robustness. Given this intricacy, the paper would be substantially strengthened by meticulous ablation studies to justify each component. For instance, a thorough analysis of the model's performance when stages (such as 1 or 4) are simplified or loss terms are removed would help distinguish critical components from those offering marginal gains and would better motivate the final design.\n\nThis substantial methodological complexity is not yet fully justified by the provided experimental validation. The cost of the method would be more compelling if benchmarked against a wider array of competitive alternatives and shown to have outstanding results. Currently, most of the results are compared almost exclusively against a single baseline (single-branch SBM). While this demonstrates that a branched model is better at modeling branches than a non-branched one, it doesn't contextualize the method within the broader field. Furthermore, Despite being motivated by cell differentiation, the paper notably omits a comparison against any of the many existing, often simpler, single-cell trajectory inference algorithms, some of which also explicitly designed to deal with branching. This makes it difficult to assess the practical utility of BranchSBM over established methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FPNW63UiSM", "reviewer": "Reviewer_2wLh", "summary": "This paper presents Branched Schrödinger Bridge Matching (BranchSBM) for learning in a new problem setting: branched Schrödinger bridges. This framework is more general than the standard approach and has the potential to be applied to certain tasks involving multiple paths. Theoretically, the problem originates from branched conditional stochastic optimal control, where there can be $K$ joint couplings ($K+1$ multi-directional branches) to represent a holistic solution from the initial condition. The authors present a bridge matching algorithm and show its efficacy on cell dynamics data.", "questions": "* Could you summarize how BranchSBM is fundamentally different from single-branch SBM with multi-modal marginals?\n* What is the general training time for single-branch SBM vs. BranchSBM for challenging problems?", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "Based on my reading, the method appears to have the following strengths.\n* The paper is relatively straightforward to understand, even though the problem itself and the surrounding materials are complex.\n* The proposed method seems to be sound (I did not check the entire appendix).\n* The authors performed experiments on biological data and showed practical aspects of BranchSBM .", "weaknesses": "I believe this work, in its current form, lacks clarity in many regards.\n\n* It seems that the branched Schrödinger bridge problem is a subset of the well-established multi-marginal Schrödinger bridge (optimal transport) problem. The claim seems to be that BranchSBM clearly specifies the notion of a source and branches, thus enabling more efficient algorithms, such as interpolant optimization. However, the manuscript is written in a way that does not clearly reveal this technical motivation and high level implications. Therefore, I think the corresponding contribution is somewhat vaguely written.\n* The overall experiments are limited, and it is not clear why single-cell data provides sufficient verification for the proposed framework. The authors are encouraged to put more effort into experiments on other high-dimensional data.\n* There is room for theoretical proof or analysis to demonstrate why the proposed scheme is fundamentally better or more efficient than single-branch SBM.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "V2v1uurcXY", "reviewer": "Reviewer_1Hyc", "summary": "The paper introduces Branched Schrödinger Bridge Matching (BranchSBM), which solves a branched generalized SB by decomposing it into a sum of unbalanced conditional SOC objectives. It learns branch-specific drifts and growth rates with mass-conservation constraints via a multi-stage procedure (interpolant → flow matching → growth → joint). On LiDAR navigation, cell differentiation, and drug-perturbation gene expression, BranchSBM improves endpoint/intermediate reconstruction over single-branch SBM.", "questions": "1. Could the method be extended to infer branching structure automatically from snapshot data? If not, please discuss how sensitive training is to incorrect or coarse endpoint clustering.\n2. The four-stage pipeline (interpolant → flow matching → growth → joint) is complex, and the contribution of each stage to the final outcome remains unclear. Could the authors report ablations on key metrics (e.g., W₁/W₂/MMD) to quantify how each stage affects trajectory fidelity and stability?\n3. The evaluation mainly contrasts with a single-branch SBM. Including other competitive baselines (e.g., DeepRUOT, SF2M, OT-CFM, MetricFM) would better contextualize where BranchSBM provides concrete advantages.\n4. Since $V_t$ guides intermediate trajectories, could the authors explore how varying its definition affects results? Currently, its influence on branching time or trajectory stability is not quantified.\n5. The perturbation experiments rely on PCA (50, 100, 150 PCs) and evaluate W distances only on the top-2 PCs. Could the authors test whether performance holds in the full PCA space to validate?\n6. The paper notes increased space complexity due to per-branch networks, but no measurements of training time or memory usage with respect to dimensionality or number of branches are reported. Including empirical scaling results (e.g., runtime vs. K or dimension) would clarify practical feasibility.\n7. Since BranchSBM explicitly learns branch-specific growth rates $g_{t,k}$, it would be informative to visualize these trajectories on biological datasets—e.g., the mouse hematopoiesis dataset. Showing how growth evolves along different branches could clarify the biological interpretability and verify whether the model captures expected proliferation or decay patterns.\n\nReferences\n1. Neklyudov, Kirill, et al. \"A computational framework for solving wasserstein lagrangian flows.\" ICML 2024.\n2. Zhang, Zhenyi et al. \"Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport.\" ICLR.\n3. Tong, Alexander, et al. \"Improving and generalizing flow-based generative models with minibatch optimal transport.\" TMLR.\n4. Tong, Alexander, et al. \"Simulation-free Schrödinger bridges via score and flow matching.\" AISTATS.\n5. Kapusniak, Kacper, et al. \"Metric flow matching for smooth interpolations on the data manifold. NeurIPS 2024.\n\nI would be happy to hear the authors’ thoughts and clarifications in the rebuttal and may adjust my score accordingly.\n\nThe reviewer wrote this review. An LLM was used only for language refinement.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1.  Frames branched GSB as a sum of Unbalanced CondSOC tasks; gives branch-wise drifts and growth with mass conservation. \n2.  Consistent W1/W2/MMD improvements vs. single-branch SBM; scales to ≥3 branches in perturbation tasks.", "weaknesses": "1. Requires branch priors/endpoint clustering. Training assumes access to clustered terminal distributions $\\pi_{1,k}$ (and uses them directly in losses), rather than discovering branches from snapshots.\n2. The work decomposes branchGSB into a sum of unbalanced CondSOC subproblems with branch-specific drifts/growth and a task-specific state cost guiding intermediate trajectories. I think several components are not entirely new. In per branch, the formulation considered here is close to unbalanced SB/RUOT except for the task-specific state cost. The neural interpolant function considered here and used for flow matching is also presented in (Neklyudov et al.,  ICML 2024). And I believe the paper misses some references. The paper discusses the connections with other methods mainly in the appendix, but it would benefit from a more explicit discussion of some key references in the main text.\n3. Not totally simulation-free in training. Although drift learning avoids state-SDE simulation (Stages 1–2), growth learning still requires time integration of $g_{t,k}$ to obtain the weight.\n4. Training complexity and the contribution of each stage remain unclear. The paper reports Stage-3 vs. Stage-4 loss trends (Table 6) but does not present ablation on metrics (e.g., W/MMD) to quantify each stage’s effect on final performance, stability, or sample quality. More granular ablations could clarify the necessity/sufficiency of each stage.\n5. Baselines are limited. Comparisons are primarily to a single-branch SBM baseline; no comparisons to other unbalanced SB/RUOT (e.g., DeepRUOT), or flow-matching methods (e.g., OT-CFM, SF2M, MetricFM) are provided, which weakens the case that BranchSBM is uniquely required.\n6. Pipeline choices may inject bias. Sensitivity to the number of branches K, and to endpoint clustering protocols, is not analyzed. \n7. Dependence on a task-specific state costs $V_t$. The approach follows intermediate trajectories “governed by a state cost”; results highlight energy/mass evolution but do not systematically study how $V_t$ choices affect accuracy, stability, or branching time.  \n8. Perturbation experiments are conducted on top-50 PCs but metrics on top-2 PCs for W distances; it is unclear why the authors choose to do so.\n9. Scalability with many branches/memory. The authors note higher space complexity (multiple branch networks), though time is said to be comparable, and inference is cheaper. Empirical scaling with large K or long horizons is not demonstrated. Also lacks analysis of training time or memory scaling with dimensionality or number of branches.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)