Conceptio › Archive › arXiv CS

arXiv CSopen access

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

neural-networks

machine learning, deep learning, neural networks

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

Junda Ying 1 Yuxuan Wang 2 Bowen Yang 3 Peijie Zhou 4 5 6 7 Lei Zhang 1 3 4 5

1. Introduction

Abstract

Recovering continuous underlying dynamics from snapshot observations is a critical challenge in single-cell biology. Due to the destructive nature of single-cell sequencing protocols, methods must infer temporal evolution without access to longitudinal trajectories of individual cells. While Optimal Transport (OT) (Kantorovich, 1958; Benamou & Brenier, 2000) has become a standard framework for this task, traditional OT seeks deterministic ODE flows between balanced probability distributions. Consequently, it fails to capture the inherent stochasticity and unbalanced nature of cellular processes driven by high biological noise and cell proliferation/apoptosis.

arXiv:2605.00545v1 [cs.LG] 1 May 2026

Inferring cellular trajectories from destructive snapshots is complicated by the challenges of stochasticity and non-conservative mass dynamics such as cell proliferation and apoptosis. Existing unbalanced Optimal Transport (OT) methods treat mass as a continuous fluid, performing inference at the population level. However, this macroscopic view often fails to capture the discrete, jump-like nature of birth-death events at single-cell resolution, which is essential for understanding lineage branching and fate decisions. We present Unbalanced Schrödinger Bridge (USB), a simulation-free framework for learning underlying dynamics that effectively integrates both stochastic and unbalanced effects which also models the discrete, jump-like birth–death dynamics at single-cell resolution. Theoretically, USB provides a tractable solution to the Branching Schrödinger Bridge (BSB) problem, offering a rigorous microscopic interpretation where individual cells undergo both Brownian motion and discrete birth-death jumps. Technically, the method implements an efficient solver by introducing a simulation-free training objective that effectively scales to high-dimensional omics data. Empirically, we demonstrate on both simulated and realworld datasets that USB not only achieves trajectory reconstruction performance better than or comparable to deterministic baselines but also uniquely enables realistic discrete simulation of birth-death dynamics at single-cell resolution.

The Schrödinger Bridge (SB) problem (Schrödinger, 1932; Léonard, 2014) is employed to model this stochasticity. It models cellular dynamics as SDEs connecting two balanced probability distributions, rather than ODEs, thereby explicitly introducing stochasticity. Parallel to the need for stochastic modeling, addressing varying cell numbers is essential (Sha et al., 2024). Both standard OT and SB assume mass conservation between time points, a condition rarely met in proliferating biological systems. To address this issue, researchers have proposed various extensions of OT for two unbalanced measures (Eyring et al., 2024; Wang et al., 2025; Peng et al., 2026). While effective at handling unbalanced population between snapshots, these standard OT-based methods often lack the capability to naturally incorporate stochasticity. To simultaneously model stochasticity and unbalanced effects, approaches based on dynamic Regularized Unbalanced Optimal Transport (RUOT) (Zhang et al., 2025) or Schrödinger Bridge with coffin states (Pariset et al., 2023) have been developed. While theoretically capable of modeling both aspects, these methods rely on computationally costly NeuralODE (Chen et al., 2018) or Iterative Proportional Fitting (IPF), rendering them computationally challenging for large data.

Beijing International Center for Mathematical Research, Peking University 2 Center for Data Science, Peking University 3 School of Mathematical Sciences, Peking University 4 Center for Quantitative Biology, Peking University 5 Center for Machine Learning Research, Peking University 6 National Engineering Laboratory for Big Data Analysis and Applications, Beijing 7 AI for Science Institute, Beijing. Correspondence to: Lei Zhang <[email protected]>, Peijie Zhou <[email protected]>.

To scale dynamical modeling to large data, recent advances have successfully introduced simulation-free paradigms– flow matching (Lipman et al., 2023)–for solving standard OT problems, along with their stochastic or unbalanced ex-

Preprint. May 4, 2026.

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

tensions (Tong et al., 2024a;b; Peng et al., 2026). These methods have rendered dynamical modeling efficient and scalable. However, there currently exists no unified framework that is simultaneously stochastic, unbalanced, and simulation-free.

Albergo & Vanden-Eijnden, 2023) is a simulation-free generative framework for learning deterministic ODE flows, and it can be coupled with optimal transport (OT) for efficiently learning dynamics (Tong et al., 2024a; Klein et al., 2024; Rohbeck et al., 2025). To accommodate more complex dynamics, prior works have introduced score-based stochastic extensions (Sohl-Dickstein et al., 2015; Song & Ermon, 2019; 2020; Song et al., 2021; Ho et al., 2020; Winkler et al., 2023; Dhariwal & Nichol, 2021; Tong et al., 2024b; Tang et al., 2025; Lee et al., 2025), unbalanced extensions handling unbalanced marginals (Eyring et al., 2024; Cao et al., 2025; Corso et al., 2025; Wang et al., 2025; Peng et al., 2026), and other generalizations (Kapuśniak et al., 2024; Zhang et al., 2024b; Atanackovic et al., 2025; Petrović et al., 2025). However, existing matching-type methods do not jointly model stochasticity and unbalanced marginals. USB introduces an unbalanced score matching framework that bridges this gap.

Furthermore, we highlight a critical limitation in current methods modeling unbalanced mass: they uniformly treat mass as a continuously varying quantity. While natural within frameworks like flow matching or NeuralODE, this assumption overlooks a fundamental characteristic of cellular proliferation and apoptosis: cell numbers are discrete values that change through jumps. To address these challenges, we present Unbalanced Schrödinger Bridge (USB), a simulation-free framework for learning underlying dynamics that accounts for both stochastic and unbalanced effects, and explicitly permits discrete birth-death dynamic simulations. Our contributions are summarized as follows:

Single-cell trajectory inference in unbalance and stochastic setting. Single-cell trajectory inference in the multi-time points setting has been tackled with NeuralODEs (Tong et al., 2020; Huguet et al., 2022; Zhang et al., 2024a; Sha et al., 2024; Gu et al., 2025; Choi & Choi, 2024) and optimal transport (OT) (Schiebinger et al., 2019; Klein et al., 2025; Halmos et al., 2025; Banerjee et al., 2025), together with their stochastic (Neklyudov et al., 2023; 2024; Albergo et al., 2025; Zhu et al., 2024; Maddu et al., 2025; Yeo et al., 2021; Chizat et al., 2022; Lavenant et al., 2024; Shi et al., 2023; Koshizuka & Sato, 2023; Bunne et al., 2023a; Chen et al., 2022; Jiang & Wan, 2024; Zhang et al., 2025; Sun et al., 2025) and unbalanced variants (Neklyudov et al., 2023; 2024; Eyring et al., 2024; Wang et al., 2025; Peng et al., 2026). Some also used branching SDEs (Ventre et al., 2024), but require lineage trees information. Flow matching can be applied to these approaches to improve scalability and stability (Tong et al., 2024a; Rohbeck et al., 2025; Klein et al., 2024; Tong et al., 2024b; Lee et al., 2025; Kapuśniak et al., 2024; Atanackovic et al., 2025). However, the field still lacks a simulation-free matching framework that jointly models stochasticity, unbalance and the discrete birth–death dynamics which need no priors.

• We propose USB, a novel framework that unifies the stochasticity with the unbalance through the lens of Branching Schrödinger Bridge problem (BSB). • We address the computational bottleneck of stochastic unbalanced modeling by developing a general simulation-free unbalanced score matching framework. • We demonstrate that USB consistently recovers groundtruth dynamics in complex landscapes and provides discrete, single-cell resolution birth-death simulations of proliferation and apoptosis.

2. Related works SB and unbalanced extensions. Many numerical algorithms have been developed to solve Schrödinger Bridge (Schrödinger, 1932) problem (SB) (De Bortoli et al., 2021; Shi et al., 2023; Bunne et al., 2023a;b; Kim et al., 2025; Wang et al., 2021; Tong et al., 2024b; Peluchetti, 2025; Somnath et al., 2023; Gushchin et al., 2024; Garg et al., 2024; De Bortoli et al., 2024; Shen et al., 2025; Noble et al., 2023). Common paradigms include converting SB into a regularized optimal transport (ROT) problem (Föllmer, 1988; Léonard, 2014; Pavon et al., 2021; Tong et al., 2024b), or stochastic optimal control (SOC) problem (Chen et al., 2016; 2022; Liu et al., 2024; Tang et al., 2025). Recently, unbalanced extensions based on coffin state (Chen et al., 2025; Pariset et al., 2023) or branching process (Baradat & Lavenant, 2021) have also been proposed, but efficient simulation-free solvers are still lacking. USB is a simulation-free solver based on the latter.

3. Preliminaries Setup. Inspired by single-cell dynamics inference, consider two unnormalized measures µ0 (x), µ1 (x), also denoted as µ0 , µ1 , at time t = 0 and t = 1 respectively defined over X ⊆ Rd . Let M+ (X ) represent the set of all absolutely continuous finite measures supported on X . The total mass of µ0 and µ1 may be different. We aim to learn the most likely stochastic process bridging µ0 and µ1 . To characterize changes in mass, two modeling paradigms are commonly used:

Flow matching and score matching. Flow matching (Lipman et al., 2023; Liu et al., 2023; Pooladian et al., 2023; 2

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

{ρt }t∈[0,1] (Benamou & Brenier, 2000). Particles are transported by the ODE flow generated by ut without growth. The equation (2) reduces to continuity equation ∂t ρt + ∇x · (ut ρt ) = 0. By introducing flow matching (Lipman et al., 2023), recent work (Tong et al., 2024a) solves the dynamic OT in a simulation-free manner. It parameterized a neural network vθ (t, x) to approximate the velocity field by minimizing a regression loss LFM (θ) = Et,x∼ρt ∥vθ (t, x) − ut (x)∥22 .

• The first employs weighted particles, assigning each particle a continuously varying mass weight and describing the temporal evolution of the measure with a Fokker–Planck equation that includes a source term; we refer to this as a continuous measure flow. • The second employs branching processes, using particle division and death to simulate jump-like changes in mass, with the particle count modeling discrete mass. For single-cell birth–death dynamics, it offers superior biological interpretability.

One key observation is that though the marginal probability path {ρt } is intractable, one can introduce tractable conditional paths ρt (·|z) and conditional velocity field ut (·|z) w.r.t some conditioning variable z. The minimization problem is equivalent to minimizing a conditional version of the loss LCFM (θ) = Et,z,x∼ρt (·|z) ∥vθ (t, x) − ut (x|z)∥22 .

Continuous measure flows. Taking both sthocastic and unbalanced effect into account, consider a time-dependent measure path ρ : Rd ×[0, 1] → R+ , a time dependent vector field u : Rd × [0, 1] → Rd , and a time dependent growth rate g : Rd × [0, 1] → R. The measure path is generated by a SDE (

dxt = ut (xt )dt + νdWt d ln mt = gt (xt )

One can choose z = (x0 , x1 ) drawn from the static OT (Kantorovich, 1958) coupling γ(x0 , x1 ) between µ0 and µ1 , and use conditional Gaussian path N (tx1 + (1 − t)x0 , ν 2 I) for determining ut (·|z). The resulting flow recovers the dynamic OT flow.

(1)

where Wt is the standard Brownian motion in Rd , and ν ∈ R is the diffusion parameter. The measure path satisfies the Fokker-Planck equation with source term ∂t ρt (x) + ∇x · (ut (x)ρt (x)) =

We point out that the algorithm is consist of two important part – the coupling and the conditional path. One can design specific coupling and conditional path to approximate other flows instead of dynamic OT flow.

ν2 ∆x ρt (x) + gt (x)ρt (x) 2 (2)

3.2. Unbalanced effect model with WFR metric To interpolate unbalanced source and target, previous works (Chizat et al., 2018a;b; Liero et al., 2018) defined WFR metric as

Branching Brownian Motion. Branching Brownian motion (BBM) is a prototypical example of branching processes. It can be described by a doublet (ν, q) where ν ∈ R is the diffusion parameter, q ∈ M+ (N) is an unnormalized measure supported on natural numbers with finite total measure called branching mechanism. The total measure P of branching mechanism λ = k∈N qk is called branching rate, and the normalized branching mechanism p = q/λ is called offspring distribution. Under BBM, a particle is equipped with an exponential clock of rate λ. It evolves according to Brownian motion with diffusion parameter ν until time τ ∼ Exp(λ). At time τ , the particle branches into k ∼ p particles. (k = 0 means that the particle vanishes). After branching, the k new particles undergo BBM independently. Restricting particles to split into two or die makes BBM well aligned with the microscopic dynamics of cell division and apoptosis. Its allowance for mass jumps also makes it a natural choice for single-cell modeling.

WFR2δ (µ0 , µ1 ) = Z 1Z 1 inf (∥u(x, t)∥22 + δ 2 ∥g(x, t)∥22 )ρt (x)dxdt ρ,g,u 0 X 2 s.t. ∂t ρ + ∇x · (ρu) = ρg, ρ0 = µ0 , ρ1 = µ1 , (3) This minimization problem is called dynamic WFR. Similar to dynamic OT, it also has a static form Z WFR2δ (µ0 , µ1 ) = 2δ 2 { inf γ0 (x, y) + γ1 (x, y) (γ0 ,γ1 )

p ∥x − y∥2 −2 γ0 (x, y)γ1 (x, y) cos( ) dxdy} 2δ (4) 2 R where R(γ0 , γ1 ) ∈ M+ (X 2 ) : X γ0 (x, y)dy = µ0 (x), X γ1 (x, y)dx = µ1 (y) is called the semicoupling. It is an unbalanced extension of the OT coupling. For fixed pair (x, y), γ0 (x, y) represents the mass at the initial time sent from x, while γ1 (x, y) represents the corresponding mass at the final time received by y. cos(x) = cos(min{x, π2 }).

3.1. Dynamic OT via Flow Matching Without stochasticity and unbalanced mass, dynamic OT provides a principled framework for interpolating between two probability measures by W2 geodesic 3

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

where µ0 and µ1 are unnomarlized measures, and Q is a BBM.

Based on these results, (Peng et al., 2026) designed a flow matching scheme to solve the problem above. They use static WFR coupling as the coupling and travelling Dirac as the conditional path. The resulting flow recovers the WFR geodesic between the source and target.

The evolution equation of BBM. According to (Baradat & Lavenant, 2021), the evolution equation of BBM is ∂t ρt =

3.3. Stochastic effect model with Schrodinger Bridge.

arg min

KL(P∥Q)

The relation to RUOT. The main result of (Baradat & Lavenant, 2021) is that the BSB problem (8) is ill-posed, and the RUOT problem

(5)

P:p0 =µ0 ,p1 =µ1

where P is a stochastic process with marginals denoted as pt . A usual choice for Q is νW, i.e. the standard Brownian motion with diffusion parameter ν. The resulting Schrödinger bridge is known as diffusion Schrödinger bridge (DSB) (De Bortoli et al., 2021; Bunne et al., 2023a; Shi et al., 2023). Following (Föllmer, 1988; Léonard, 2014), one can disintegrate the KL-divergence into two part. Z KL(P∥Q) = KL(P01 ∥Q01 )+ KL(Pxy ∥Qxy )P01 (dx, dy) (6) where P01 , Q01 are the marginal of P, Q at time 0, 1, and Pxy , Qxy are the conditional process of P, Q given start point x and end point y. The second term measured the difference of conditional path. It can be minimized to 0 by choosing Pxy = Qxy . Thus, it is sufficient to only minimize the first term ⋆

P =

arg min KL(P01 ∥Q01 )

(9)

P where r = k̸=1 (k − 1)qk . In a weak sense, it is equivalent to evolve according to Brownian motion with diffusion parameter ν in position, and to grow exponentially in mass.

To allow stochastic dynamics bridging two probability measures, (Schrödinger, 1932) proposed the Schrödinger bridge problem. It asks to find a most likely stochastic process P⋆ bridging normalized µ0 and µ1 w.r.t a reference stochastic process Q. P⋆ =

ν2 ∆x ρt + rρt 2

RUOT(µ0 , µ1 ) = Z 1Z 1 ∥u(x, t)∥22 + Ψ(g) ρt (x)dxdt inf ρ,g,u 0 X 2 ν2 s.t.∂t ρ + ∇x · (ρu) = ∆x ρ + ρg, ρ0 = µ0 , ρ1 = µ1 , 2 (10) is a relaxation of it, i.e. the dual of (8) and (10) happens to be the same. The growth penalization Ψ has a complicated form dependent on ν and q. When p0 = p2 = 12 and pk = 0, k ̸= 0, 2, i.e. particles split into two or die with no preference, Ψ is determined by its Legendre transform g Ψ∗ν,λ (g) = νλ(cosh( ) − 1) (11) λ These results connect continuous measure flow with branching processes, allowing us to realize BSB in continuous measure flow framework.

4. Simulation-free training of USB problem

(7)

In this section, we first establish a general framework for learning unbalanced stochastic dynamics (4.1,4.2,4.3,4.4), and then focus on the specific case of USB (4.5,4.6).

p0 =µ0 ,p1 =µ1 ,

It is also known as the static Schrödinger bridge. When Q is νW with the same source measure µ0 , it reduces to a regularized OT (ROT) problem. Utilizing these results, (Tong et al., 2024b) proposed SF2 M, a simulation-free framework for solving diffusion Schrödinger bridge. It uses ROT coupling as the coupling, and uses (νW)xy (which is Brownian bridge) as conditional path to recover the Schrödinger bridge between µ0 and µ1 .

4.1. Unbalanced score matching loss design Given a vector field ut and a rate function gt that generate the measure path ρt by (2), we can view the Fokker-Planck equation with source term as a continuity equation with source term ∂t ρt (x) + ∇x · (u◦t (x)ρt (x)) = gt (x)ρt (x)

3.4. Unbalanced Schrodinger Bridge

(12)

where u◦t (x) = ut (x) − ν2 ∇x ln ρt (x). It is called the drift of probability flow ODE (Tong et al., 2024b) under deterministic settings. The term st (x) = ∇x ln ρt (x) is known as the score function. The equation (12) can be generated by a measure flow ODE instead of SDE ( dxt = u◦t (xt )dt (13) d ln mt = gt (xt )dt

To take both stochastic effect and unbalanced effect into account, (Baradat & Lavenant, 2021) proposed an unbalanced formulation of Schrödinger bridge. They replace the reference process with branching Brownian motion (BBM) to allow growth. The BSB problem is then defined as P⋆ = arg min KL(P∥Q) (8) P:p0 =µ0 ,p1 =µ1

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

Theorem 4.1. The marginal vector field and rate function (18) generates the marginal measure path (17) for any q(z) independent of x and t. The score function and u◦ also satisfies the marginalization reR lation st (x) = st (x|z) ρt (x|z)q(z) dz, u◦t (x) = ρt (x) R ◦ ut (x|z) ρt (x|z)q(z) dz. ρt (x)

and the original SDE (1) generating (2) can be viewed as  ν2  dxt = (u◦t (xt ) + st (xt ))dt + νdWt (14) 2  d ln mt = gt (xt )dt Thus, to recover the unbalanced stochastic dynamics (1), it is sufficient to learn a time-dependent vector field vθ (x, t), a time-dependent growth rate gθ (x, t), and a time-dependent score function sθ (x, t) parametrized by neural networks, with loss function specified by minimizing the intractable unbalanced score matching objective (USM)

Conditional Gaussian measure path. As a convenient instance of conditional path, conditional Gaussian mesure path (CGMP) is introduced ρ̃t (x|z) = N (x|ηt (z), σt2 (z)I)

The conditional vector field, rate function, and score function of CGMP are easy to compute. (Details in Appendix A.2) Proposition 4.2. For CGMP (19),  ′ u◦ (x|z) = σt (z) x − η (z) + η ′ (z)   t t t  σt (z)   gt (x|z) = ∂t ln mt (z) (20)    x − ηt (z)   st (x|z) = − σt2 (z)

LUSM (θ) = Z 1Z 2 2 (∥vθ (x, t) − u◦t (x)∥2 + ∥gθ (x, t) − gt (x)∥2 0

+ λ2 (t) ∥sθ (x, t) − st (x)∥2 )ρt (x)dxdt (15) We utilized the weight for score loss λ(t) adopted from (Tong et al., 2024b) for numerical stability. The details are left to Appendix B.1. 4.2. Conditional path construction

4.3. Conditional Loss design

Conditional measure path. Following the approach of defining a conditional measure path in analogy to unbalanced flow matching (Peng et al., 2026), we define a conditional measure path w.r.t the condition variable z such that ρt (x|z) = mt (z)ρ̃t (x|z) where the time-dependent conditional measure ρt (x|z) is decoupled into a time dependent mass mt (z) and a time-dependent conditional probability density ρ̃t (x|z). The conditional velocity field, growth rate satisfy the conditional Fokker-Planck equation with source term ∂t ρt (x|z) + ∇x · (ut (x|z)ρt (x|z)) = (16) ν2 ∆x ρt (x|z) + ρt (x|z)gt (x|z) 2

We can regress the conditionals which are tractable by minimizing the conditional unbalanced score matching objective (CUSM) LCUSM (θ) = 2

Et∼U [0,1],z∼q(z),x∼ρ̃t (x|z) mt (z)(∥vθ (x, t) − u◦t (x|z)∥2 2

The following theorem recovers the classical results of matching algorithms that one can minimize the intractable marginal objective (15) by minimizing the tractable conditional objective (21). The proof is left to Appendix A.3. Theorem 4.3. If ρt (x) > 0 for all x ∈ X and t ∈ [0, 1], and q(z) is independent of x and t, then LUSM (θ) = LCUSM (θ) + C, where C is independent of θ. Thus they have identical gradients w.r.t θ, i.e.

Marginal measure path. Assuming z ∼ q(z), we define the marginal measure path from the conditional measure path Z ρt (x|z)q(z)dz

as well as the marginal vector field and growth rate Z   ut (x) = ut (x|z) ρt (x|z)q(z) dz  ρt (x) Z   gt (x) = gt (x|z) ρt (x|z)q(z) dz ρt (x)

+ ∥gθ (x, t) − gt (x|z)∥2 + λ2 (t) ∥sθ (x, t) − st (x|z)∥2 ) (21) Here we weight the regression loss by mass mt (z) to deal with unbalanced mass. In a deterministic setting, the loss reduces to the CUFM loss proposed by (Peng et al., 2026). In a balanced setting where mt (z) ≡ 1, gt (x|z) ≡ 0, (21) naturally reduces to the standard conditional score matching loss (Tong et al., 2024b).

and we define the conditional score function as st (x|z) = ∇x ln ρt (x|z).

ρt (x) =

(19)

(17)

∇θ LUSM (θ) = ∇θ LCUSM (θ). (18) 4.4. A general framework for learning unbalanced stochastic dynamics

The following marginalization theorem connects the conditionals and marginals. The proof is left to Appendix A.1.

In sections above, we have established a general framework for learning unbalanced stochastic dynamics. Once q(z) and 5

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots Table 1. Examples of coupling and conditional path construction Algorithm

coupling

OT-CFM SF2 M WFR-FM USB

OT coupling π ROT coupling π2ν 2 WFR semi-coupling (γ0 , γ1 ) RUOT semi-coupling (γ0 , γ1 )

tx1 + (1 − t)x0 N (tx1 + (1 − t)x0 , ν 2 t(1 − t)I) Rt ωds

1 1 At2 + Bt + C t m1−t 0 m1

0 As2 +Bs+C

N (tx1 + (1 − t)x0 , ν 2 t(1 − t)I)

is easy to obtain  1 − 2t   u◦t (x|z) = x − ((1 − t)x0 + tx1 ) + (x1 − x0 )   t(1 − t)   gt (x|z) = ln m1 (x0 , x1 ) − ln m0 (x0 , x1 )    (1 − t)x0 + tx1 − x   st (x|z) = ν 2 t(1 − t) (24)

the conditionals u◦t (x|z), gt (x|z), st (x|z) are specified, we can regress the marginals u◦t (x), gt (x), st (x) by minimizing (21). A common choice for z is z = (x0 , x1 ) i.e. the pair of start point and end point from the source and target respectively. In this case, q(z) stands for some coupling between the source and target, while the conditionals stand for the conditional path between two Diracs m0 δx0 , m1 δx1 . m0 , m1 are determined by the semi-coupling, which reduces to one coupling in balanced cases, resulting in m0 = m1 = 1. We list the semi-coupling/coupling and conditional path used by some previous methods, as well as USB in Table 1. In the following part of this section, we will focus on the conditional path and semi-coupling of USB.

4.6. Static coupling construction for USB The minimization of the first term of (22) results in a static USB semi-coupling (γ0 , γ1 ). Though hard to obtain, note that RUOT is a relaxation of USB (Baradat & Lavenant, 2021). Thus, we use the semi-coupling induced by the corresponding RUOT problem (10) to approximate the static USB semi-coupling. Here we restrict the branching mechanism of the referencing BBM to p0 = p2 = 12 , pk = 0, k ̸= 0, 2 to imitate the real cell division and apoptosis.

4.5. Conditional path construction for USB We disintegrate the KL-divergence (8) Z KL(P∥Q) = KL(P01 ∥Q01 )+ KL(Pxy ∥Qxy )P01 (dx, dy) (22) where Q is now a BBM. The second term can be minimized to 0 by choosing Pxy = Qxy . Note that it is hard to condition on a BBM in strong sense due to its branching nature. But, the conditional measure path can be constructed based on the evolution equation of BBM (9). The solution of (9) is a weighted Gaussian ρt (x) = ert N (x|0, ν 2 I) where the position x follows a Brownian motion with diffusion parameter ν, and the mass varies linearly in log-scale. Intuitively, since BBM branching events follow an exponential distribution and result in a multiplication of the particle count, the particle number can be viewed approximately as a Poisson process in log-scale. The log-linear mass can be viewed as a limit of Poisson process in log-scale with infinitesimal increment. Thus, given a pair of Dirac (m0 δx0 , m1 δx1 ) as source and target, we bridge them by a Poisson-Brownian bridge, i.e. we bridge (x0 , x1 ) with Brownian bridge, and bridge the mass with linear interpolation in log-scale, which is the limit of Poission bridge (Conforti et al., 2015) with infinitesimal increment. The details are presented in Appendix B.2.  dxt = x1 − xt dt + νdWt 1−t (23)  d ln mt = (ln m1 − ln m0 )dt

In practice, it is also hard to calculate the semi-coupling for general RUOT problem due to the complexity of the growth penalty Ψ. For efficiency, we approximate the RUOT problem with a WFR problem (3), which is easy to solve, by expanding Ψ to second order. More computational details are presented in Appendix B.3. We also discuss the main difficulty in Appendix E.

5. USB workflow for trajectory inference Multi-time points USB. Given samples from unbalanced measures µi at K + 1 discrete time points, t = t0 , t1 , . . . , tK , multi-time points USB tends to find a most likely evolution bridging these marginal measures. P⋆ =

arg min

KL(P∥Q)

(25)

P:pi =µi ,i=0,1,··· ,K+1

where Q is a BBM. The KL-disintegration yields KL(P∥Q) = KL(P{0,1,··· ,K} ∥Q{0,1,··· ,K} )+ Z KL(Px0 ···xK ∥Qx0 ···xK )P{0,1,··· ,K} (dx0 , · · · dxK ) (26) Similar to the two-time points case, the second term can be minimized to 0 by choosing Px0 ···xK = Qx0 ···xK . Due to Markov property, conditioning on multiple time points is equivalent to conditioning on consecutive pairs and stitching the corresponding conditional measure paths together.

As a CGMP ρt (x|z) = mt (z)N (ηt (z), ν 2 t(1−t)I) where mt (z) = m0 (z)1−t m1 (z)t , ηt (z) = (1 − t)x0 + tx1 , z = (x0 , x1 ), the conditionals of Poisson-Brownian bridge

The minimization of the first term results in K pairs of semicouplings. These semi-couplings can be approximated by 6

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots Table 2. Comparison of trajectory inference algorithms based on Fokker-Planck-like equation(2).

solving a multi-marginal RUOT problem RUOT(µ0 , · · · , µK ) = Z Z tK − t0 tK 1 inf ∥u(x, t)∥22 + Ψ(g) ρt (x)dxdt ρ,g,u 2 t0 X 2 ν2 ∆x ρ + ρg, ρi = µi , i = 0, · · · , K 2 (27) Due to Markov property again, the dynamics between different consecutive time pairs are independent. According to (Peng et al., 2026), under this independence assumption, the multi-time problem reduces to solving RUOT between successive time points. The proof is left to Appendix A.4.

s.t.∂t ρ + ∇x · (ρu) =

Proposition 5.1. The solution to the multi-time RUOT problem (27) is equivalent to the concatenation of the solutions of successive time points.

ν2 sθ (xt , t))dt + νdWt 2  d ln mt = gθ (xt , t)dt

stochastic

Simulation-free

discrete birth-dearth dynamics

✗ ✗ ✓ ✓ ✗ ✓

✗ ✓ ✗ ✓ ✓ ✓

✓ ✓ ✓ ✗ ✗ ✓

✗ ✗ ✗ ✗ ✗ ✓

Table 3. Mean W1 and RME (only for unbalanced methods) on synthetic datasets. For the methods that exhibit randomness in inference, we report the mean value and standard deviation over 5 runs. Best results are in bold, and the second best are underlined. Method

Simulation (2D)

Dyngen (5D)

Gaussian (1000D)

W1 (↓)

RME (↓)

W1 (↓)

RME (↓)

W1 (↓)

RME (↓)

0.298 0.311 0.224±0.007 0.148 0.474 0.045 0.043±0.002 0.079±0.003 0.093 0.046 0.019 0.019±0.000

— — — — — 0.014 0.017±0.001 0.008±0.002 0.010 0.006 0.001 0.002±0.000

1.371 1.767 1.277±0.017 0.965 1.415 0.512 0.623±0.032 0.522±0.008 1.204 0.598 0.135 0.131±0.001

— — — — — 0.047 0.065±0.011 0.177±0.007 0.097 0.037 0.005 0.007±0.000

2.833 3.794 3.543±0.002 2.858 3.438 2.263 3.785±0.009 2.813±0.004 2.771 3.010 2.233 2.136±0.002

— — — — — 0.127 0.303±0.070 0.041±0.006 0.033 0.037 0.044 0.004±0.004

6. Experiment results

Inference schemes. We provide two inference modes for USB, namely, Continuous Inference and Branching Inference. For Continuous Inference, USB simulates the trajectory by following SDE dxt = (vθ (xt , t) +

unbalanced

OT-CFM SF2 M WFR-FM DeepRUOT BranchSBM USB

MMFM Metric FM SF2M MIOFlow BranchSBM TIGON DeepRUOT Var-RUOT UOT-FM VGFM WFR-FM USB

Training workflow. In practice, the neural networks share the parameters across different pairs of time points. For convenient, we choose the condition variable z = (x0 .x1 ) ∼ 0 .x1 ) γ0 (x0 .x1 ), and set m0 (z) = 1, m1 (z) = γγ01 (x (x0 .x1 ) for each two-time points cases. We described the training workflow of USB on multi-time points in Appendix F.

 

Algorithm

USB accurately bridges µt0 , · · · , µtK . We evaluate how accurate can USB mathces the measure at observed time points on three synthetic datasets: Simulation Gene (Simulation), Dyngen and the 1000D Gaussian Mixtures (Gaussian). The distribution-matching accuracy is measured by the 1Wasserstein distance (W1 ) between the normalized true measure and the normalized USB simulated measure, and the mass-matching accuracy is measured by Relative Mass Error (RME), the relative error of predicted total mass (Experiment details in Appendix C.2). USB demonstrates superior performance across all three datasets, achieving best accuracy on distribution-matching task, and best or second-best accuracy on mass-matching task, performing on par with the top-performing baseline, while consistently and significantly outperforming its balanced counterpart, SF2 M (Tong et al., 2024b) (Table 3). We also evaluate USB on several real datasets in Appendix C.

(28)

The continuous mode aims to match the source and target, and recover the USB dynamics on population-level, while the Branching Inference aims to simulate the discrete single-cell birth-death dynamics at single-cell resolution. We start from a cell at x0 whose position evolves according to the upper SDE of (28). We also simulates a non-homogeneous branching process with rate |gθ (xt , t)| to decide when the cell undergoes branching. At the branching event, the cell divides into two (gθ > 0) or dies (gθ < 0). After branching, all living cells follow the simulation process above independently. More details are presented in Appendix B.4. Pseudocodes are presented in Appendix F.

USB accurately interpolates the unobserved time points. We evaluate how well can USB recover the true underlying dynamics by hold-one-out experiments. For a datasets with K + 1 time points t0 , t1 , · · · , tK , holding out one time point ti , 1 ≤ i ≤ K − 1, USB is trained on the remaining K time points, and the W1 is evaluated at the holding out time point ti . The mean performance on all holding out time points are presented in Table 4. Across EMT (Cook & Vanderhyden, 2020), EB (Moon et al., 2019), CITE-seq (Lance et al., 2022), and mouse hematopoiesis data (Mouse) (Weinreb et al., 2020), USB achieves best interpolation

A full model for single-cell dynamics. We discussed several trajectory inference algorithms related to USB in Appendix D. Among algorithms based on Fokker-Plancklike equations (2), USB models both the stochastic ef2 fect ν2 ∆x ρt (x) and the unbalanced effect gt (x)ρt (x) in a simulation-free manner, which also allows discrete simulation for birth-death dynamics (Table 2). 7

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots Table 4. Mean W1 over held-out time points on EMT, EB, CITE and Mouse datasets. For the methods that exhibit randomness in inference, we report the mean value and standard deviation over 5 runs. Best results are in bold, and the second best are underlined. Method

EMT (10D)

EB (50D)

CITE (50D)

Mouse (50D)

MMFM Metric FM SF2M MIOFlow BranchSBM TIGON DeepRUOT Var-RUOT UOT-FM VGFM WFR-FM USB

0.323 0.314 0.308±0.001 0.325 0.369 0.360 0.323±0.002 0.320±0.003 0.322 0.301 0.298 0.298±0.0001

11.213 10.726 10.986±0.006 10.960 11.988 11.080 10.075±0.004 11.035±0.017 11.344 10.370 10.157 10.177±0.0001

38.521 37.342 38.333±0.002 39.574 39.125 38.159 37.892±0.002 38.393±0.029 38.649 37.386 37.221 37.087±0.0001

8.263 7.753 8.646±0.004 7.779 8.586 6.868 6.847±0.003 8.672±0.040 9.332 8.496 6.586 6.988±0.0001

Figure 2. Single-cell resolution birth-death dynamics simulation

tions from the same cell to generate 8 trajectories (ν = 0.1) using Branching Inference, distinguished by different colors (Fig.2). The total mass of the Dyngen data initially decreases and subsequently increases, evolving into two imbalanced branches, with the lower branch containing a higher cell count than the upper one. We observe that cell fates diverse significantly even when originating from the same initial state: the majority of cells undergo early apoptosis, while some trajectories traverse towards distinct branches, experiencing subsequent division or death events. These results demonstrate USB’s capacity to model not only cell division but also apoptosis and complex bifurcation behaviors.

Figure 1. Learned growth rate on the Gaussian 1000D dataset. Left panel: WFR-FM; Right panel: USB

accuracy on EMT and CITE, while performs comparable with top baselines on others. This demonstrates that USB successfully captures the underlying cellular dynamics. USB recovers the underlying birth-death dynamics. To evaluate how well can USB recover the underling birthdeath dynamics, we calculated the Pearson correlation ( X22 gcorr ) between the true growth rate g = αg 1+X 2 and the 2 predicted growth rate on the Simulation Gene dataset. Averaging on 5 runs, USB gets a mean Pearson correlation of 0.9739, showing high consistency to the ground truth. The growth rate plot on the 1000D Gaussian dataset also shows that USB recovers a more plausible growth rate than WFR, which is designed for capturing unbalanced effect (Figure 1). In the 1000D Gaussian settings, the upper cluster expands from 100 to 1,000 cells without displacement, whereas the lower cluster maintains its total counts (400), bifurcating into two 200-cell clusters. Consequently, the ground-truth growth rate is high in the upper cluster and zero across the lower clusters; USB recovers this behavior more faithfully than WFR. The difference is because of the different underlying dynamic assumptions of the two algorithms. More detailed discussion is presented in Appendix D.2.

7. Conclusion, limitation and discussion In this work, we introduced USB, a framework rooted in the BSB problem. By employing unbalanced score matching, USB efficiently integrates the stochastic and unbalanced effect inherent in single-cell dynamics, and explicitly permits the discrete simulation of birth-death dynamics. We validated the effectiveness and robustness of the method on both synthetic and real scRNA-seq datasets, demonstrating its significant potential for biological trajectory inference. A limitation of USB is that the semi-coupling of BSB and its RUOT relaxation are intractable. Consequently, we use WFR for approximation. The underlying mathematics merits further investigation. However, the unbalanced score matching framework we have established is indeed general: provided that the static semi-coupling and the conditional path of the target problem, our approach can be extended to other domains. This generality renders USB applicable to a wider range of dynamical modeling or machine learning scenarios involving stochasticity and unbalance, or featuring microscopic branching and discrete mass variations.

USB simulates discrete birth-death dynamics at singlecell resolution. On the Dyngen dataset, we initiated simula8

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

Impact Statements

Bunne, C., Stark, S. G., Gut, G., et al. Learning singlecell perturbation responses using neural optimal transport. Nature methods, 20:1759–1768, 2023b.

This paper presents work whose goal is to advance the field of machine learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.

Cannoodt, R., Saelens, W., Deconinck, L., and Saeys, Y. Spearheading future omics analyses using dyngen, a multi-modal simulator of single cells. Nature communications, 12:3942, 2021.

Acknowledgments

Cao, Z., Zhong, Y., and Deng, L.-J. Taming flow matching with unbalanced optimal transport into fast pansharpening, 2025.

This work was supported by the National Natural Science Foundation of China (NSFC No. 12225102 to L.Z., 8206100646 to P.Z., T2321001 to P.Z. & L.Z., and 12288101 to P.Z. & L.Z.) and the National Key R&D Program of China No.2024YFA0919500 to L.Z. We thank the anonymous referees for their valuable feedback and constructive suggestions.

References

Chen, R. T. Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. In Bengio, S., Wallach, H., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.

Albergo, M., Boffi, N. M., and Vanden-Eijnden, E. Stochastic interpolants: A unifying framework for flows and diffusions. Journal of Machine Learning Research, 26 (209):1–80, 2025.

Chen, T., Liu, G.-H., and Theodorou, E. Likelihood training of schrödinger bridge using forward-backward SDEs theory. In International Conference on Learning Representations, 2022.

Albergo, M. S. and Vanden-Eijnden, E. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, 2023.

Chen, Y., Georgiou, T., and Pavon, M. On the relation between optimal transport and schrödinger bridges: A stochastic control viewpoint. Journal of Optimization Theory and Applications, 169:671–691, 2016.

Atanackovic, L., Zhang, X., Amos, B., Blanchette, M., Lee, L. J., Bengio, Y., Tong, A., and Neklyudov, K. Meta flow matching: Integrating vector fields on the wasserstein manifold, 2025.

Chen, Y., Georgiou, T. T., and Pavon, M. Optimal survival strategies for diffusive flows: A schrödinger bridge approach to unbalanced transport. SIAM Review, 67(3): 579–604, 2025. doi: 10.1137/25M176581X.

Banerjee, A., Lee, H., Sharon, N., and Moosmüller, C. Efficient trajectory inference in wasserstein space using consecutive averaging. In Li, Y., Mandt, S., Agrawal, S., and Khan, E. (eds.), Proceedings of The 28th International Conference on Artificial Intelligence and Statistics, volume 258 of Proceedings of Machine Learning Research, pp. 2260–2268. PMLR, 03–05 May 2025.

Chizat, L., Peyré, G., Schmitzer, B., and Vialard, F.-X. Scaling algorithms for unbalanced transport problems, 2017. Chizat, L., Peyré, G., Schmitzer, B., and Vialard, F.-X. An interpolating distance between optimal transport and fisher–rao metrics. Foundations of Computational Mathematics, 18(1):1–44, 2018a.

Baradat, A. and Lavenant, H. Regularized unbalanced optimal transport as entropy minimization with respect to branching brownian motion, 2021.

Chizat, L., Peyré, G., Schmitzer, B., and Vialard, F.-X. Unbalanced optimal transport: Dynamic and kantorovich formulations. Journal of Functional Analysis, 274(11): 3090–3123, 2018b.

Benamou, J.-D. and Brenier, Y. A computational fluid mechanics solution to the monge-kantorovich mass transfer problem. Numerische Mathematik, 84(3):375–393, 2000.

Chizat, L., Zhang, S., Heitz, M., and Schiebinger, G. Trajectory inference via mean-field langevin in path space. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 16731–16742. Curran Associates, Inc., 2022.

Bunne, C., Hsieh, Y.-P., Cuturi, M., and Krause, A. The schrödinger bridge between gaussian measures has a closed form. In Ruiz, F., Dy, J., and van de Meent, J.-W. (eds.), Proceedings of The 26th International Conference on Artificial Intelligence and Statistics, volume 206 of Proceedings of Machine Learning Research, pp. 5802– 5833. PMLR, 25–27 Apr 2023a.

Choi, J. and Choi, J. Scalable simulation-free entropic unbalanced optimal transport, 2024. 9

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

Conforti, G., Léonard, C., Murr, R., and Rœlly, S. Bridges of Markov counting processes. Reciprocal classes and duality formulas. Electronic Communications in Probability, 20(none):1 – 12, 2015. doi: 10.1214/ECP.v20-3697.

Garg, J., Zhang, X., and Zhou, Q. Soft-constrained Schrödinger bridge: a stochastic control approach. In Dasgupta, S., Mandt, S., and Li, Y. (eds.), Proceedings of The 27th International Conference on Artificial Intelligence and Statistics, volume 238 of Proceedings of Machine Learning Research, pp. 4429–4437. PMLR, 02– 04 May 2024.

Cook, D. P. and Vanderhyden, B. C. Context specificity of the emt transcriptional response. Nature communications, 11:2142, 2020.

Gu, A., Chien, E., and Greenewald, K. Partially observed trajectory inference using optimal transport and a dynamics prior. In The Thirteenth International Conference on Learning Representations, 2025.

Corso, G., Somnath, V. R., Getz, N., Barzilay, R., Jaakkola, T., and Krause, A. Composing unbalanced flows for flexible docking and relaxation. In The Thirteenth International Conference on Learning Representations, 2025.

Gushchin, N., Kholkin, S., Burnaev, E., and Korotin, A. Light and optimal schrödinger bridge matching. In Forty-first International Conference on Machine Learning, 2024.

De Bortoli, V., Thornton, J., Heng, J., and Doucet, A. Diffusion schrödinger bridge with applications to score-based generative modeling. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 17695–17709. Curran Associates, Inc., 2021.

Halmos, P., Liu, X., Gold, J., Chen, F., Ding, L., and Raphael, B. J. Dest-ot: Alignment of spatiotemporal transcriptomics data. Cell Systems, 16(2), 2025.

De Bortoli, V., Korshunova, I., Mnih, A., and Doucet, A. Schrödinger bridge flow for unpaired data translation. In Globerson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomczak, J., and Zhang, C. (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 103384–103441. Curran Associates, Inc., 2024. doi: 10.52202/079017-3285.

Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 6840– 6851. Curran Associates, Inc., 2020. Huguet, G., Magruder, D. S., Tong, A., Fasina, O., Kuchroo, M., Wolf, G., and Krishnaswamy, S. Manifold interpolating optimal-transport flows for trajectory inference. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 29705–29718. Curran Associates, Inc., 2022.

Dhariwal, P. and Nichol, A. Diffusion models beat gans on image synthesis. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 8780–8794. Curran Associates, Inc., 2021. Eyring, L., Klein, D., Uscidda, T., Palla, G., Kilbertus, N., Akata, Z., and Theis, F. J. Unbalancedness in neural monge maps improves unpaired domain translation. In The Twelfth International Conference on Learning Representations, 2024.

Jiang, Q. and Wan, L. A physics-informed neural sde network for learning cellular dynamics from time-series scrna-seq data. Bioinformatics, 40(Supplement 2):ii120– ii127, 09 2024. ISSN 1367-4811. doi: 10.1093/ bioinformatics/btae400.

Fatras, K., Zine, Y., Majewski, S., Flamary, R., Gribonval, R., and Courty, N. Minibatch optimal transport distances; analysis and applications, 2021.

Kantorovich, L. V. On the translocation of masses. Management Science, 5(1):1–4, 1958. ISSN 00251909, 15265501.

Flamary, R., Courty, N., Gramfort, A., Alaya, M. Z., Boisbunon, A., Chambon, S., Chapel, L., Corenflos, A., Fatras, K., Fournier, N., et al. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1–8, 2021.

Kapuśniak, K., Potaptchik, P., Reu, T., Zhang, L., Tong, A., Bronstein, M., Bose, A. J., and Di Giovanni, F. Metric flow matching for smooth interpolations on the data manifold. In Globerson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomczak, J., and Zhang, C. (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 135011–135042. Curran Associates, Inc., 2024. doi: 10.52202/079017-4291.

Föllmer, H. Random fields and diffusion processes. In Hennequin, P.-L. (ed.), École d’Été de Probabilités de Saint-Flour XV–XVII, 1985–87, pp. 101–203, Berlin, Heidelberg, 1988. Springer Berlin Heidelberg. ISBN 978-3540-46042-8.

Kim, J. H., Kim, S., Moon, S., Kim, H., Woo, J., and Kim, W. Y. Discrete diffusion schrödinger bridge matching for 10

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

graph transformation. In The Thirteenth International Conference on Learning Representations, 2025.

Liu, X., Gong, C., and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2023.

Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In ICLR (Poster), 2015.

Léonard, C. A survey of the schrödinger problem and some of its connections with optimal transport. Discrete and Continuous Dynamical Systems, 34(4):1533–1574, 2014. ISSN 1078-0947. doi: 10.3934/dcds.2014.34.1533.

Klein, D., Uscidda, T., Theis, F., and Cuturi, M. Genot: Entropic (gromov) wasserstein flow matching with applications to single-cell genomics. In Globerson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomczak, J., and Zhang, C. (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 103897–103944. Curran Associates, Inc., 2024. doi: 10.52202/079017-3301.

Maddu, S., Chardès, V., and Shelley, M. J. Inferring biological processes with intrinsic noise from cross-sectional data, 2025. Moon, K. R., Van Dijk, D., Wang, Z., Gigante, S., Burkhardt, D. B., Chen, W. S., Yim, K., Elzen, A. v. d., Hirn, M. J., Coifman, R. R., et al. Visualizing structure and transitions in high-dimensional biological data. Nature biotechnology, 37:1482–1492, 2019.

Klein, D., Palla, G., and Lange, M. o. Mapping cells through time and space with moscot. Nature, 638:1065–1075, 2025. Koshizuka, T. and Sato, I. Neural lagrangian schrödinger bridge: Diffusion modeling for population dynamics. In The Eleventh International Conference on Learning Representations, 2023.

Neklyudov, K., Brekelmans, R., Severo, D., and Makhzani, A. Action matching: Learning stochastic dynamics from samples. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 25858–25889. PMLR, 23–29 Jul 2023.

Lance, C., Luecken, M. D., Burkhardt, D. B., Cannoodt, R., Rautenstrauch, P., Laddach, A., Ubingazhibov, A., Cao, Z.-J., Deng, K., Khan, S., Liu, Q., Russkikh, N., Ryazantsev, G., Ohler, U., data integration competition participants, N. . M., Pisco, A. O., Bloom, J., Krishnaswamy, S., and Theis, F. J. Multimodal single cell data integration challenge: Results and lessons learned. In Kiela, D., Ciccone, M., and Caputo, B. (eds.), Proceedings of the NeurIPS 2021 Competitions and Demonstrations Track, volume 176 of Proceedings of Machine Learning Research, pp. 162–176. PMLR, 06–14 Dec 2022.

Neklyudov, K., Brekelmans, R., Tong, A., Atanackovic, L., Liu, Q., and Makhzani, A. A computational framework for solving Wasserstein lagrangian flows. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 37461–37485. PMLR, 21–27 Jul 2024.

Lavenant, H., Zhang, S., Kim, Y.-H., and Schiebinger, G. Toward a mathematical theory of trajectory inference. The Annals of Applied Probability, 34(1A):428–500, 2024.

Noble, M., De Bortoli, V., Doucet, A., and Durmus, A. Treebased diffusion schrödinger bridge with applications to wasserstein barycenters. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 55193–55236. Curran Associates, Inc., 2023.

Lee, J., Moradijamei, B., and Shakeri, H. Multi-marginal stochastic flow matching for high-dimensional snapshot data at irregular time points, 2025. Liero, M., Mielke, A., and Savaré, G. Optimal entropytransport problems and a new hellinger–kantorovich distance between positive measures. Inventiones mathematicae, 211:969–1117, 2018.

Pariset, M., Hsieh, Y.-P., Bunne, C., Krause, A., and Bortoli, V. D. Unbalanced diffusion schrödinger bridge, 2023. Pavon, M., Trigila, G., and Tabak, E. G. The data-driven schrödinger bridge. Communications on Pure and Applied Mathematics, 74(7):1545–1573, 2021.

Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, 2023.

Peluchetti, S. BM$ˆ2$: Coupled schrödinger bridge matching. Transactions on Machine Learning Research, 2025. ISSN 2835-8856.

Liu, G.-H., Lipman, Y., Nickel, M., Karrer, B., Theodorou, E., and Chen, R. T. Q. Generalized schrödinger bridge matching. In The Twelfth International Conference on Learning Representations, 2024.

Peng, Q., Wang, Z., Ying, J., Sun, Y., Nie, Q., Zhang, L., Li, T., and Zhou, P. Wfr-fm: Simulation-free dynamic unbalanced optimal transport, 2026. 11

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

Petrović, K., Atanackovic, L., Moro, V., Kapuśniak, K., İsmail İlkan Ceylan, Bronstein, M., Bose, A. J., and Tong, A. Curly flow matching for learning non-gradient field dynamics, 2025.

Somnath, V. R., Pariset, M., Hsieh, Y.-P., Martinez, M. R., Krause, A., and Bunne, C. Aligned diffusion Schrödinger bridges. In Evans, R. J. and Shpitser, I. (eds.), Proceedings of the Thirty-Ninth Conference on Uncertainty in Artificial Intelligence, volume 216 of Proceedings of Machine Learning Research, pp. 1985–1995. PMLR, 31 Jul– 04 Aug 2023.

Pooladian, A.-A., Ben-Hamu, H., Domingo-Enrich, C., Amos, B., Lipman, Y., and Chen, R. T. Q. Multisample flow matching: Straightening flows with minibatch couplings. In ICML, pp. 28100–28127, 2023.

Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.

Pra, P. D. A stochastic control approach to reciprocal diffusion processes. Appl Math Optim, 23:313–329, 1 1991. doi: https://doi.org/10.1007/BF01442404. Rohbeck, M., Bunne, C., Brouwer, E. D., Huetter, J.-C., Biton, A., Chen, K. Y., Regev, A., and Lopez, R. Modeling complex system dynamics with flow matching across time and conditions. In The Thirteenth International Conference on Learning Representations, 2025.

Song, Y. and Ermon, S. Improved techniques for training score-based generative models. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 12438–12448. Curran Associates, Inc., 2020.

Schiebinger, G., Shu, J., Tabaka, M., Cleary, B., Subramanian, V., Solomon, A., Gould, J., Liu, S., Lin, S., Berube, P., et al. Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 176(4):928–943, 2019.

Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.

Schrödinger, E. Sur la théorie relativiste de l’électron et l’interprétation de la mécanique quantique. Annales de l’Institut Henri Poincaré, 2(4):269–310, 1932.

Sun, Y., Zhang, Z., Wang, Z., Li, T., and Zhou, P. Variational regularized unbalanced optimal transport: Single network, least action, 2025.

Sha, Y., Qiu, Y., Zhou, P., and Nie, Q. Reconstructing growth and dynamic trajectories from single-cell transcriptomics data. Nature Machine Intelligence, 6(1):25– 39, 2024.

Tang, S., Zhang, Y., Tong, A., and Chatterjee, P. Branched schrödinger bridge matching. In The Exploration in AI Today Workshop at ICML 2025, 2025. Theodoropoulos, P., Saravanos, A. D., Theodorou, E. A., and Liu, G.-H. Momentum multi-marginal schrödinger bridge matching, 2025.

Shen, Y., Berlinghieri, R., and Broderick, T. Multi-marginal schrödinger bridges with iterative reference refinement. In Li, Y., Mandt, S., Agrawal, S., and Khan, E. (eds.), Proceedings of The 28th International Conference on Artificial Intelligence and Statistics, volume 258 of Proceedings of Machine Learning Research, pp. 3817–3825. PMLR, 03–05 May 2025.

Tong, A., Huang, J., Wolf, G., Van Dijk, D., and Krishnaswamy, S. TrajectoryNet: A dynamic optimal transport network for modeling cellular dynamics. In III, H. D. and Singh, A. (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 9526–9536. PMLR, 13–18 Jul 2020.

Shi, Y., De Bortoli, V., Campbell, A., and Doucet, A. Diffusion schrödinger bridge matching. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 62183–62223. Curran Associates, Inc., 2023.

Tong, A., FATRAS, K., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Wolf, G., and Bengio, Y. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, pp. 1–34, 2024a. ISSN 2835-8856.

Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequilibrium thermodynamics. In Bach, F. and Blei, D. (eds.), Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 2256–2265, Lille, France, 07–09 Jul 2015. PMLR.

Tong, A. Y., Malkin, N., Fatras, K., Atanackovic, L., Zhang, Y., Huguet, G., Wolf, G., and Bengio, Y. Simulation-free Schrödinger bridges via score and flow matching. In Dasgupta, S., Mandt, S., and Li, Y. (eds.), Proceedings of The 27th International Conference on Artificial Intelligence 12

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

and Statistics, volume 238 of Proceedings of Machine Learning Research, pp. 1279–1287. PMLR, 02–04 May 2024b. Ventre, E., Forrow, A., Gadhiwala, N., Chakraborty, P., Angel, O., and Schiebinger, G. Trajectory inference for a branching sde model of cell differentiation, 2024. Wang, D., Jiang, Y., Zhang, Z., Gu, X., Zhou, P., and Sun, J. Joint velocity-growth flow matching for single-cell dynamics modeling. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. Wang, G., Jiao, Y., Xu, Q., Wang, Y., and Yang, C. Deep generative learning via schrödinger bridge. In Meila, M. and Zhang, T. (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 10794– 10804. PMLR, 18–24 Jul 2021. Weinreb, C., Rodriguez-Fraticelli, A., Camargo, F. D., and Klein, A. M. Lineage tracing on transcriptional landscapes links state to fate during differentiation. Science, 367(6479):eaaw3381, 2020. Winkler, L., Ojeda, C., and Opper, M. A score-based approach for training schrödinger bridges for data modelling. Entropy, 25(2), 2023. ISSN 1099-4300. doi: 10.3390/e25020316. Yeo, G. H. T., Saksena, S. D., and Gifford, D. K. Generative modeling of single-cell time series with prescient enables prediction of cell trajectories with interventions. Nature communications, 12:3222, 2021. Zhang, J., Larschan, E., Bigness, J., and Singh, R. scNODE: generative model for temporal single cell transcriptomic data prediction. Bioinformatics, 40(Supplement 2):ii146– ii154, 09 2024a. ISSN 1367-4811. Zhang, X., Pu, Y., Kawamura, Y., Loza, A., Bengio, Y., Shung, D. L., and Tong, A. Trajectory flow matching with applications to clinical time series modelling. In Globerson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomczak, J., and Zhang, C. (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 107198– 107224. Curran Associates, Inc., 2024b. doi: 10.52202/ 079017-3404. Zhang, Z., Li, T., and Zhou, P. Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport. In The Thirteenth International Conference on Learning Representations, 2025. Zhu, Q., Zhao, B., Zhang, J., Li, P., and Lin, W. Governing equation discovery of a complex system from snapshots, 2024. 13

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

A. Proofs Proofs of theorems and propositions A.1. Proof of Theorem 4.1 Theorem 4.1. The marginal vector field and rate function (18) generates the marginal measure path (17) for any q(z) independent of x and t. The score function and u◦ also satisfies the marginalization relation st (x) = R R dz, u◦t (x) = u◦t (x|z) ρt (x|z)q(z) dz. st (x|z) ρt (x|z)q(z) ρt (x) ρt (x) Proof. To show that the marginals generate the marginal measure path (17), it is sufficient to show that the marginal measure ρt (x) satisfies the Fokker-Planck equation with source term (2). ∂t ρt (x) =

d dt Z

Z ρt (x|z)q(z)dz ∂t ρt (x|z)q(z)dz

In (1), we use the Fokker-Planck equation with source term of the conditional measure ρt (x|z) (16). In (2), we use the definition of marginal vector field and rate function (18). Thus, the marginal measure ρt (x) satisfies the Fokker-Planck equation with source term (2). By direct calculation, st (x) = ∇x ln ρt (x) 1 ∇x ρt (x) = ρt (x) Z 1 ∇x ρt (x|z)q(z) dz = ρt (x) Z 1 = ∇x ρt (x|z)q(z) dz ρt (x) Z 1 = ∇x ln ρt (x|z)ρt (x|z)q(z) dz ρt (x) Z ρt (x|z)q(z) = st (x|z) dz ρt (x) ν2 u◦t (x) = ut (x) − st (x) 2 Z Z ν2 = ut (x|z) ρt (x|z)q(z) dz − st (x|z) ρt (x|z)q(z) dz ρt (x) ρt (x) 2 Z ν2 = ut (x|z) − st (x|z) ρt (x|z)q(z) dz ρt (x) 2 Z = u◦t (x|z) ρt (x|z)q(z) dz ρt (x) The relation between conditionals and marginals is verified. The proof is completed. 14

□

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

A.2. Proof or Proposition 4.2 Proposition 4.2. For CGMP (19),  σt′ (z)  ◦  x − ηt (z) + ηt′ (z) ut (x|z) =  σt (z)   gt (x|z) = ∂t ln mt (z)    x − ηt (z)   st (x|z) = − σt2 (z) Proof. The CMGP satisfies the following conditional Fokker-Planck equation with source term (16) ∂t ρt (x|z) + ∇x · (ut (x|z)ρt (x|z)) =

ν2 ∆x ρt (x|z) + ρt (x|z)gt (x|z) 2

Expanding the conditional measure into mass and density, the equation become mt (z)∂t ρ̃t (x|z)+∂t mt (z)ρ̃t (x|z)+mt (z)∇x ·(ut (x|z)ρ̃t (x|z)) =

ν2 mt (z)∆x ρ̃t (x|z)+mt (z)ρ̃t (x|z)gt (x|z) (29) 2

Since mt (z) > 0 for all t, we devide the both sides of the equation by mt (z) and get ∂t ρ̃t (x|z) + ∂t ln mt (z)ρ̃t (x|z) + ∇x · (ut (x|z)ρ̃t (x|z)) =

ν2 ∆x ρ̃t (x|z) + ρ̃t (x|z)gt (x|z) 2

(30)

We further reorganize it to ∂t ρ̃t (x|z) + ∇x · (ut (x|z)ρ̃t (x|z)) =

ν2 ∆x ρ̃t (x|z) + ρ̃t (x|z) gt (x|z) − ∂t ln mt (z) 2

(31)

ν2 ∆x ρ̃t (x|z) 2

(32) 2

(33)

Note that this equation is exactly the continuity equation of the conditional Gaussian path in conditional flow matching. Thus, the conditional vector field of CGMP shares the same form with the conditional vector field of conditional Gaussian σ ′ (z) path in balanced conditional flow matching, which is u◦t (x|z) = σtt (z) x − ηt (z) + ηt′ (z). The conditional score function of CGMP is just the score function of a family of Gaussian distributions st (x|z) = ∇x ln ρt (x|z) (1)

= ∇x ln ρ̃t (x|z)

(2)

= ∇x ln N (x|ηt (z), σt2 (z)I)

=−

x − ηt (z) σt2 (z)

In (1), note that the mass term mt (z) is independent of x, thus the score function of ρt and ρ̃t are equal. In (2), we use the definition of CGMP ρ̃t (x|z) = N (x|ηt (z), σt2 (z)I) (19). □

The proof is completed. 15

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

A.3. Proof of Theorem 4.3 Theorem 4.3. If ρt (x) > 0 for all x ∈ X and t ∈ [0, 1], and q(z) is independent of x and t, then LUSM (θ) = LCUSM (θ) + C, where C is independent of θ. Thus they have identical gradients w.r.t θ, i.e. ∇θ LUSM (θ) = ∇θ LCUSM (θ). Proof. Recall the two objectives Z 1Z 2 2 2 LUSM (θ) = (∥vθ (x, t) − u◦t (x)∥2 + ∥gθ (x, t) − gt (x)∥2 + λ2 (t) ∥sθ (x, t) − st (x)∥2 )ρt (x)dxdt 0

LCUSM (θ) = Et∼U [0,1],z∼q(z),x∼ρ̃t (x|z) mt (z)(∥vθ (x, t) − u◦t (x|z)∥2 + ∥gθ (x, t) − gt (x|z)∥2 + λ2 (t) ∥sθ (x, t) − st (x|z)∥2 ) By direct calculation, Z

∥vθ (x, t)∥22 − 2⟨vθ (x, t), u◦t (x)⟩ + ∥gθ (x, t)∥22 − 2⟨gθ (x, t), gt (x)⟩ + λ2 (t) ∥sθ (x, t)∥22 − 2⟨sθ (x, t), st (x)⟩ ρt (x)dx + C. Z Z (1) = Et∼U [0,1] ∥vθ (x, t)∥22 − 2⟨vθ (x, t), u◦t (x|z)⟩ + ∥gθ (x, t)∥22 − 2⟨gθ (x, t), gt (x|z)⟩ X + λ2 (t) ∥sθ (x, t)∥22 − 2⟨sθ (x, t), st (x|z)⟩ mt (z)ρ̃t (x|z)q(z) dz dx + C Z Z 2 2 = Et∼U [0,1] ∥vθ (x, t) − u◦t (x|z)∥2 + ∥gθ (x, t) − gt (x|z)∥2 X 2 + λ2 (t) ∥sθ (x, t) − st (x|z)∥2 mt (z)ρ̃t (x|z)q(z)dz dx + C 2 2 = Et∼U [0,1],z∼q(z),x∼ρ̃t (x|z) mt (z) ∥vθ (x, t) − u◦t (x|z)∥2 + ∥gθ (x, t) − gt (x|z)∥2 2 + λ2 (t) ∥sθ (x, t) − st (x|z)∥2 + C

LUSM (θ) = Et∼U [0,1]

= LCUSM (θ) + C where C is independent of θ. In (1), we use the relation between conditionals and marginals (18, 4.1). Since the constant C is independent of θ, we have ∇θ LUSM (θ) = ∇θ LCUSM (θ) □

The proof is completed. A.4. Proof of Proposition 5.1

Proposition 5.1. The solution to the multi-time RUOT problem (27) is equivalent to the concatenation of the solutions of successive time points. Proof. (adopted from (Peng et al., 2026) with modification) Let t0 < t1 < · · · < tK be the discrete time points. For each k = 1, . . . , K, denote by (ρ(k) (x, t), u(k) (x, t), g (k) (x, t)) the optimal solution of the two-time RUOT problem between µtk−1 and µtk : Z tk Z tk − tk−1 RUOT(µtk−1 , µtk ) = inf ∥u(x, t)∥22 + Ψ(g(x, t)) ρ(x, t) dx dt (34) ρ,u,g t 2 X k−1 subject to ∂t ρ + ∇x · (ρu) =

ν2 ∆x ρ + ρg, 2

ρtk−1 = µtk−1 , ρtk = µtk .

Now define the concatenated trajectory ρ̃(x, t) = ρ(k) (x, t),

ũ(x, t) = u(k) (x, t), 16

g̃(x, t) = g (k) (x, t),

for t ∈ [tk−1 , tk ].

(35)

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots (k)

(k+1)

Since ρtk = µtk = ρtk

, the concatenated triple (ρ̃, ũ, g̃) is admissible for the multi-time RUOT problem.

The corresponding cost is tK − t 0 2 =

K X k=1

Z tK Z t0

∥ũ(x, t)∥22 + Ψ(g̃(x, t)) ρ̃(x, t) dx dt

(36)

tK − t 0 RUOT(µtk−1 , µtk ). tk − tk−1

Hence RUOT({µt0 , . . . , µtK }) ≤

K X tK − t0 k=1

tk − tk−1

RUOT(µtk−1 , µtk ).

(37)

To see that the equality must hold, assume by contradiction that there exists another admissible solution (ρ′ , u′ , g ′ ) such that Z Z K X tK − t0 tK tK − t 0 ∥u′ (x, t)∥22 + Ψ(g ′ (x, t)) ρ′ (x, t) dx dt < RUOT(µtk−1 , µtk ). 2 t k − tk−1 t0 X k=1

Since the total cost is a sum over the disjoint intervals [tk−1 , tk ], this strict inequality implies that there must exist at least one index k ∗ such that Z Z tK − t0 tk∗ tK − t 0 ∥u′ (x, t)∥22 + Ψ(g ′ (x, t)) ρ′ (x, t) dx dt < RUOT(µtk∗ −1 , µtk∗ ). ∗ − tk ∗ −1 2 t k tk∗ −1 X 0 gives Dividing both sides by the positive factor tk∗tK−t−t k∗ −1 Z Z tk∗ − tk∗ −1 tk∗ ∥u′ (x, t)∥22 + Ψ(g ′ (x, t)) ρ′ (x, t) dx dt < RUOT(µtk∗ −1 , µtk∗ ). 2 tk∗ −1 X

But this contradicts the definition of RUOT(µtk∗ −1 , µtk∗ ) as the minimal cost between µtk∗ −1 and µtk∗ . Thus, no admissible solution can have strictly smaller cost than the concatenated one. We conclude that the concatenated trajectory (ρ∗ (x, t), u∗ (x, t), g ∗ (x, t)) =

K [

(ρ(k) , u(k) , g (k) ),

t ∈ [tk−1 , tk ],

(38)

k=1

achieves the infimum of the multi-time problem and is therefore optimal. The proof is completed.

□

Remark: This is true only when the dynamics between adjacent time pairs (tk−1 , tk ) are independent. In multi-marginal methods, such as MMFM (Rohbeck et al., 2025), 3MSBM (Theodoropoulos et al., 2025) and MMSFM (Lee et al., 2025) where high order continuity or global connections between different time pairs are introduced, proposition (5.1) will never hold.

B. Implementation details B.1. Weighting schedule λ(t) p For numerical stability, we use the same weighting schedule λ(t) = νt(1 − t) as (Tong et al., 2024b) to normalize the score loss. With this weighting schedule, the regression target of the weighted score net λ(t)sθ (x, t) is simplified to a standard Gaussian noise ϵt ∼ N (0, I) instead of the conditional score function (24) including a division over ν 2 t(1 − t) which may cause numerical instability when t is near to 0, 1. 2

2 λ (t) ∥sθ (x, t) − st (x|z)∥2 = 2

ηt − x λ(t)sθ (x, t) − p ν t(1 − t) 2

(1)

= ∥λ(t)sθ (x, t) + ϵt ∥2

x−ηt In (1), note that x ∼ N (ηt , ν 2 t(1 − t)I), hence √ ν

t(1−t)

= ϵt ∼ N (0, I). 17

(39)

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

B.2. Poisson-Brownian bridge Since BBM is a branching process, it generates multiple end points from single start point. Thus, it is impossible to find a conditional BBM trajectory given its start point and end point. One can only try to find a conditional process in weak sense, i.e. find a conditional measure path of BBM. According to the evolution equation of BBM (9), we know that BBM undergoes diffusion spatially, while following exponential growth in mass. At a microscopic level, if we approximate that all cells share the same exponential splitting clock, then the logarithm of the total cell count is a Poisson process, the limit of which recovers this exponential growth. For numerical tractability, in what follows we determine the conditional path for this weighted particle approximation. Given two Diracs m0 δx0 , m1 δx1 , the stochastic process bridging x0 , x1 is a Brownian motion fixed on two sides, i.e. a Brownian bridge. N (tx1 + (1 − t)x0 , ν 2 t(1 − t)I) (40) The stochastic process bridging ln m0 , ln m1 is a Poisson process fixed on two sides, i.e. a Poisson bridge (Conforti et al., 2015). In cases where initial count and final count are two natural numbers N0 ≤ N1 , the Poisson bridge between them is P(t) ∼ N0 + Bernoulli(N1 − N0 , t)

(41)

Now, consider adjusting the unit of mass from 1 to ∆m, so that the counts N0 and N1 become total masses M0 = ∆mN0 and M1 = ∆mN1 , respectively. Consequently, the Poisson bridge connecting these states is P(t) ∼ M0 + ∆m · Bernoulli(N1 − N0 , t)

(42)

When N1 − N0 is large, one utilize the central limit theorem (CLT) P(t) ≈ M0 + ∆m · N ((N1 − N0 )t, (N1 − N0 )t(1 − t)) = M0 + N (∆m(N1 − N0 )t, ∆m2 (N1 − N0 )t(1 − t)) = N (tM1 + (1 − t)M0 , ∆m(M1 − M0 )t(1 − t))

(43)

= N (tM1 + (1 − t)M0 , O(∆m)) Let ∆m → 0, we finally recover the linear mass variation on log scale. P(t) ≈ N (tM1 + (1 − t)M0 , O(∆m)) → δtM1 +(1−t)M0

(∆m → 0)

(44)

We point out that this approximation is justified due to the large sample size of single-cell sequencing data. Therefore, given two Dirac measures (representing the precise states at the start and endpoints), we connect their positions using a Brownian bridge, and linearly connect their masses on log scale. We term this the Poisson-Brownian bridge (23). B.3. Approximating RUOT semi-coupling In a word, we approximate the RUOT semi-coupling by a static WFR semi-coupling which can be easily constructed from the coupling of an optimal entropy-transport problem (OET), and the OET problem can be easily solved using the POT package (Flamary et al., 2021). In order to find the semi-coupling induced by the USB problem (8), we approximate it by solving its relaxation – RUOT (10). When the referencing BBM has diffusion parameter ν, branching rate λ, and branching mechanism p0 = p2 = 21 , pk = 0, k ̸= 0, 2, the corresponding RUOT is Z Z 1 1 RUOT(µ0 , µ1 ) = inf ∥u(x, t)∥22 + Ψ(g(x, t)) ρ(x, t) dx dt ρ,u,g 2 0 X (45) ν2 ∆x ρ + ρg, ρ0 = µ0 , ρ1 = µ1 s.t. ∂t ρ + ∇x · (ρu) = 2 with growth penalty Ψν,λ (g) determined by its Legendre transform (Baradat & Lavenant, 2021) g Ψ∗ν,λ (g) = νλ(cosh( ) − 1) λ 18

(46)

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

By applying Legendre transform to Ψ⋆ , we obtain the explicit form of the growth penalty. p p Ψ(g) = 1 − 1 + g 2 + g ln( 1 + g 2 + g) g Ψν,λ (g) = νλΨ( ) λ

(47)

Since it is hard to find a static form for the RUOT problem or solve it with such a growth penalty, we approximate it by expanding the growth penalty to second order. 1 2 g + o(g 3 ) 2 ν 2 g + o(g 3 ) Ψν,λ (g) = 2λ Ψ(g) =

(48)

Doing so, we approximate the original RUOT problem with a WFR problem which is easy to solve. Z Z 1 1 ∥u(x, t)∥22 + δ 2 |g(x, t)|2 ρ(x, t) dx dt ρ,u,g 2 0 (49) Z X p ∥x − y∥2 2 = 2δ inf γ (x, y) + γ (x, y) − 2 γ (x, y)γ (x, y) cos( ) dxdy 0 1 0 1 2δ (γ0 ,γ1 )∈(M+ (X 2 ))2 X 2

RUOT(µ0 , µ1 ) ≈ WFR2δ (µ0 , µ1 ) = inf

p where δ = ν/2λ, cos(x) = cos(min{x, 0}). The second row is known as the static WFR problem which is a minimization problem respect to the semi-coupling (γ0 , γ1 ) instead of (ρ, u, g). To solve the static WFR problem, (Chizat et al., 2018b; Liero et al., 2018) proved that the static WFR problem is equivalent to the following optimal entropy-transport (OET) problem. Z ∥x − y∥2 OETδ (µ0 , µ1 ) = 2δ 2 inf 2 { −2 ln cos( )γ(x, y)dxdy 2δ γ∈M+ (X ) 2 X (50) Z Z + KL( γ(x, y)dy∥µ0 (x)) + KL( γ(x, y)dx∥µ1 (y))} X

An OET problem can be easily solved using the generalized Sinkhorn-Knopp matrix scaling algorithm (Chizat et al., 2017) which is well implemented in the POT package (Flamary et al., 2021). To construct the semi-coupling (γ0 , γ1 ) from the OET coupling γ, we follow the results of (Liero et al., 2018; Peng et al., 2026). Theorem B.1. (Peng et al., 2026) Let γ be the optimal coupling of the OET problem (50), then the semi-coupling µ (x), γ1 (x, y) = R γ(x,y) µ (y) solves the static WFR problem (4). γ0 (x, y) = R γ(x,y) γ(x,z)dz 0 γ(z,y)dz 1 X

We summarize the workflow for the calculation of the semi-coupling (γ0 , γ1 ) as following. Algorithm 1 Semi-coupling calculation Workflow Require:pSample-able distributions µ0 , µ1 , diffusion parameter ν, branching rate λ. 1: δ = ν/2λ 2: γ ← OETδ (µ0 , µ1 ) (Defined in 4) γ(x,y) 3: γ0 (x, y) ← R µ (x), γ1 (x, y) ← R γ(x,y) µ (y) (Theorem B.1) γ(x,z)dz 0 γ(z,y)dz 1 X X 4: return (γ0 , γ1 )

One can determine the triplet (ν, λ, δ) by knowing any two of them. We point out that, when fixing thep stochastic reference ν, the WFR growth penalty coefficient δ also has a intuitive meaning of growth reference. Since δ = ν/2λ, the smaller the growth penalty coefficient δ, the larger the growth reference λ, and vice versa. Thus, for convenience, we implicitly choose the branching rate λ by giving ν and δ. Remark. As mentioned above, some RUOT problem e.g. WFR do have an equivalent OET formulation which is easy to solve (Chizat et al., 2018a;b). We tried to follow their approach, but failed. A discussion on their approach and the difficulties here can be found in Appendix E. 19

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

B.4. Branching inference Branching inference aims to simulate the discrete single-cell birth-death dynamics at single-cell resolution. We start from a cell at x0 whose position evolves according to the marginal SDE. ν2 sθ (xt , t))dt + νdWt (51) 2 All SDE simulations were implemented using Euler-Maruyama scheme, and all ODE simulations were implemented using forward Euler scheme. Detailed inference workflows can be found in Appendix F. dxt = (vθ (xt , t) +

We also simulates a non-homogeneous branching process to decide when the cell divides or dies. According to the evolution equation of BBM (9), the relation between the learned growth rate gθ (x, t) and the corresponding branching rate λ, offspring distribution p at (x, t) is gθ (x, t) = λ(x, t)(p2 (x, t) − p0 (x, t)) (52) For simplicity, we assume that stochasticity in cells is solely confined to changes in gene expression, while their apoptosis or division is entirely determined by gene expression and time; that is, the offspring distribution p(x, t) is a one-hot vector given (x, t). Hence, we have  λ(x, t) = |gθ (x, t)|    p = 1 2 gθ (x,t)≥0 (53)  p = 1 0 g (x,t)<0  θ   pk = 0, k ̸= 0, 2 where 1A is the indicator function of event A. Based on the above, we linked the learned growth rates to the branching dynamics. We simulate the branching event of each cell as a non-homogeneous jump process. In details, for each SDE simulation time step ∆t, the probability of a cell branching within that step is given by p = |gθ (x, t)|∆t. Therefore, we sample a random variable α ∼ U[0, 1]. If α ≤ p, branching occurs, otherwise it does not. When branching occurs: if gθ (x, t) ≥ 0, the cell divides into two, and the resulting daughter cells are subsequently simulated independently according to these aforementioned rules; else, the current cell dies. We further point out that the mean mass increase of continuous inference and branching inference on log-scale are the same. In any time step ∆t, the log mass increment is gθ (x, t)∆t, while the mean log mass increment of branching inference is λ(x, t)∆t(p2 − p0 ) = |gθ (x, t)|∆t(1gθ (x,t)≥0 − 1gθ (x,t)<0 ) = |gθ (x, t)|∆t · sgn(gθ (x, t)) = gθ (x, t)∆t which is the same. Also, since continuous inference and branching inference follow the same SDE spatially, they generate the same measure in weak sense.

C. Additional results In this section, we provide a detailed description of the computational resources (C.1), metrics (C.2), and datasets used. We further carry out sensitivity analyses for the growth penalization δ (C.3), diffusion parameter ν (C.4), and the mini-batch OT (C.7,C.9), and assess the algorithm’s scalability with respect to the number of cells (C.9) and the dimensionality (C.7). C.1. Experimental Details The experiments were performed on a shared high-performance computing cluster with NVIDIA H800 GPU and 128 CPU cores. The architecture of the neural networks used to parameterize vθ (x, t), gθ (x, t) and sθ (x, t) are Multilayer Perceptrons with 256 hidden channels and 5 layers. We use the LeakyReLU activation function. These networks were optimized using Adam (Kingma & Ba, 2015), and implemented using Pytorch. The OT problems are solved using the Python Optimal Transport (POT) package (Flamary et al., 2021). C.2. Evaluation Metrics To evaluate model performance on measure reconstruction, we decoupled the task into two parts: distribution-matching and mass-matching. Two measures p, q are said to be the same, if and only if they have the same normalized distributions p̃, q̃, 20

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

and same total mass mp , mq , since p = mp p̃, q = mq q̃. In the followings, let p = mp p̃ be the true measure, and q = mq q̃ be the predicted measure. We use the 1-Wasserstein distance (W1 ) to measure the similarity between normalized predicted distribution q̃ and true distribution p̃. Z W1 (p̃, q̃) =

min π∈Π(p̃,q̃)

where 2

∥x − y∥2 dπ(x, y)

Π(p̃, q̃) = {π(x, y) ∈ M+ (X )|

(54)

Z π(x, y)dy = p̃(x),

π(x, y)dx = q̃(y)} X

In practice, p̃, q̃ are empirical distributions consist or weighted particles. The above minimization problem can be easily solved by linear program or Sinkhorn algorithm using POT package. We use Relative Mass Error (RME) to assess how well the model captures cell population growth. The RME is defined as RME(p, q) =

|mp − mq | mq

(55)

In previous literature such as (Sha et al., 2024; Zhang et al., 2025; Wang et al., 2025; Peng et al., 2026), RME is formulated as P | i wi (tk ) − nk /n0 | RME(tk ) = nk /n0 where wi (tk ) is the inferred weight of cell i at time tk , and nk is the number of cells at time k. One can easily check the two formulations are the same. To generate predicted measures at time tk , we start form the true measure at t0 . Each cell is assigned with a mass weight wi (0) = 1). For continuous inference, each cell travels following the learned SDE dynamics, and its weight evolves following the learned birth-death ODE. For branching inference, only the position of each cell follows the SDE, the weight stay invariant, while the cell may vanish or get a copy at the branching event. At tk , the true measure is consists of nk cells with weight 1 from the dataset, and the predicted measure is consists of weighted cells which are still alive at tk . The total mass is defined as the sum of weights (also the cell number for branching inference). For some datasets, we also perform a hold-out experiment. The model was trained on all but one time point and evaluated on that time point. All simulations were started at the initial time point t0 . To compare USB with other methods, we implement several existed methods on all datasets. For WFR-FM (Peng et al., 2026) and VGFM (Wang et al., 2025), we used their default hyperparameter settings, since the datasets evaluated in our work are also used in their work, and our model sizes are consistent. For branchSBM (Tang et al., 2025), we follow the hyperparameter settings in Table 9 of their paper. For other simulation-free methods (MMFM (Rohbeck et al., 2025), Metric FM (Kapuśniak et al., 2024), SF2M (Tong et al., 2024b), UOT-FM (Eyring et al., 2024)), we trained 3000 epochs with cosine annealing learning rate decay on high-dimensional real datasets (Mouse, Cite, EB50D, EB100D), 1000 epochs with constant learning rate on other datasets. For simulation-based methods (MIOFlow (Huguet et al., 2022), TIGON (Sha et al., 2024) and DeepRUOT (Zhang et al., 2025), Var-RUOT (Sun et al., 2025)), we follow the hyperparameter setting recommended by DeepRUOT, which is the SOTA in simulation-based methods. C.3. Simulation Gene Data The Simulation Gene data is adopted from (Zhang et al., 2025). It simulates a synthetic gene regulatory network via the following equations dX1 α1 X12 + β = − δ1 X1 + η1 ξt 2 dt 1 + α1 X1 + γ2 X22 + γ3 X32 + β dX2 α2 X22 + β = − δ2 X2 + η2 ξt dt 1 + γ1 X12 + α2 X22 + γ3 X32 + β dX3 α3 X32 = − δ3 X3 + η3 ξt dt 1 + α3 X32 21

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

where Xi represents the concentration of gene i. The model features a toggle switch system (mutual inhibition and self-activation) between X1 and X2 with inhibition by X3 and activation by an external signal β. The self-activation rate, inhibition rate, and degradation rate of each gene are denoted as αi , γi , and δi , respectively. Also, noise terms (ηi ξt ) are added to each governing ODE to introduce stochasticity. The growth is introduced by probabilistic cell division X22 with probability g = αg 1+X 2 %. Upon division, each child inherits the gene expression level of its parent, and adds an 2 independent noise ηd N (0, 1) to each gene. The data is generated by simulating the governing equations from two source. One is located at a steady states of the system with no growth, thus exhibits equilibrium during time. The other one exhibits both transition and growth. The data is recorded at five time points [0, 8, 16, 24, 32]. Sensitivity study of growth penalty δ. When referencing on a BBM with diffusion parameter ν and branching rate λ, we have to solvep a RUOT problem (10) to get the semi-coupling. We approximate it by solving a WFR problem with growth penalty δ = ν/2λ (B.3). The parameter δ explicitly represents the ratio of the strength of transition and growth. With excessively high δ, cells tend to transit rather than grow. Changes in the relative mass between the two trajectories will not be realized through population growth. Instead, it will be compensated by incorrect transfers between cell states, causing the model to learn an incorrect vector field and underlying dynamics (Fig.3 left panel). With excessively low δ, cells tend to grow rather than transit. When the measure is transported between time points, it tends to first decrease the mass and then increase it after the transfer is completed. Errors in the vector field can be compensated by the unpanelized large growth rate, causing the model to learn an incorrect vector field and dynamics (Fig.3 right panel). To balance the mutually compensating effects of transfer and growth, we chose a moderate δ = 1.3 in the main text to penalize the growth, thereby learning the correct vector field and underlying dynamics (Fig.3 middle panel). Sensitivity analysis results for different δ are shown in Table 5. Both an excessively high or low δ (δ = 0.5, 2.5) results in a sub-optimal performance, while δ in a proper range (1 to 2) exhibit comparable performance, demonstrating the robustness of USB to δ.

Figure 3. Learned trajectories on the Simulation Gene dataset (ν = 0.001). From left to right: δ = 0.5, 1.3, 2.5. Table 5. Sensitivity analysis for parameter δ on simulation gene dataset. Parameter δ = 0.5 δ = 1.0 δ = 1.3 δ = 1.5 δ = 1.7 δ = 2.0 δ = 2.5

t=1

t=2

t=3

t=4

RME

0.033±6 × 10−5 0.022±5 × 10−5 0.020±4 × 10−5 0.022±5 × 10−5 0.025±5 × 10−5 0.022±2 × 10−5 0.023±3 × 10−5

0.020±6 × 10−5 0.003±5 × 10−5 0.000±2 × 10−5 0.002±4 × 10−6 0.001±3 × 10−6 0.000±3 × 10−6 0.000±6 × 10−6

0.078±1 × 10−4 0.026±3 × 10−5 0.021±2 × 10−5 0.029±7 × 10−5 0.027±8 × 10−5 0.026±2 × 10−4 0.041±0.002

0.073±1 × 10−4 0.006±1 × 10−5 0.001±1 × 10−5 0.006±5 × 10−6 0.004±5 × 10−6 0.000±2 × 10−5 0.004±2 × 10−4

0.096±2 × 10−4 0.024±5 × 10−5 0.019±1 × 10−5 0.026±8 × 10−5 0.022±7 × 10−5 0.034±6 × 10−4 0.029±6 × 10−4

0.112±2 × 10−4 0.008±2 × 10−5 0.002±2 × 10−5 0.009±6 × 10−6 0.006±8 × 10−6 0.002±9 × 10−5 0.007±4 × 10−4

0.086±2 × 10−4 0.022±5 × 10−5 0.018±2 × 10−5 0.025±3 × 10−5 0.021±1 × 10−4 0.047±3 × 10−4 0.026±7 × 10−4

0.130±2 × 10−4 0.011±2 × 10−5 0.003±3 × 10−5 0.009±7 × 10−6 0.009±9 × 10−6 0.007±1 × 10−4 0.008±4 × 10−4

C.4. Dyngen Data We adopt the same dataset as in (Huguet et al., 2022; Wang et al., 2025). It was simulated by Dyngen (Cannoodt et al., 2021). The dataset contains 728 cells, znd was reduced to 5 dimensions using PHATE (Moon et al., 2019). The data exhibits a complex dynamics which contains both varying mass across time points and bifurcation. The bifurcation is also unbalanced in some sense that the cell number in the lower branch is larger than the upper branch. Due to the complexity, it is a challenging task to recover its dynamics. USB accurately recovers the bifurcating trajectories, and predicts a higher 22

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

growth rate in the lower branch, capturing the unbalanced bifurcation dynamics (Fig.4). In the main text, we set δ = 1.7. 4 0.88

0.73

Predicted rate

0.58

0.43

0.28

0.13

2 2

Figure 4. Learned trajectories and growth on the Dyngen dataset (δ = 1.7, ν = 0.1). Left panel: trajectories; Right panel: growth rate.

We also summarized the detailed quantitative results across different time points in Table 6. As shown, USB achieves the top-2 lowest W1 and RME in all of the time points. Sensitivity study of diffusion parameter ν. We further provide a sensitivity analysis for the diffusion parameter ν in the bottom of Table 6. As the level pf noise becomes larger, the accuracy of the measure transportation decreases. But, we point out that even for a extremely large noise (ν = 0.1), the W1 distance and RME of USB remains the second lowest in most experiments, demonstrating the robustness and power of USB. In the main text, if there is no specific notice, we set ν = 0.001. Remark. When setting ν = 0.5, the training diverges, thus ν = 0.1 is a extremely large noise level for Dyngen data. Table 6. Comparison of method performance over time on the Dyngen dataset and sensitivity analysis for parameter ν. Best results are in bold, and the second best are underlined (results of ν = 0.01, 0.05, 0.1 are not compared). Method

t=1

t=2

t=3

t=4

RME

MMFM Metric FM SF2M MIOFlow BranchSBM TIGON DeepRUOT Var-RUOT UOT-FM VGFM WFR-FM USBν = 0.001

0.574 0.892 0.637 0.420 0.926 0.446 0.454 0.315 0.652 0.335 0.110 0.109±0.0002

— — — — — 0.033 0.011 0.128 0.008 0.001 0.003 0.002±0.0000

1.704 2.347 1.266 0.640 1.171 0.584 0.481 0.548 0.780 0.312 0.098 0.093±0.0002

— — — — — 0.060 0.070 0.336 0.077 0.073 0.007 0.002±0.0001

1.499 2.030 1.415 1.537 2.081 0.415 0.870 0.630 1.252 1.109 0.211 0.180±0.0013

— — — — — 0.023 0.104 0.222 0.090 0.041 0.008 0.015±0.0004

1.706 1.799 1.790 1.263 1.481 0.603 0.688 0.593 2.130 0.634 0.121 0.142±0.0040

— — — — — 0.071 0.074 0.023 0.213 0.033 0.002 0.008±0.0008

USBν = 0.01 USBν = 0.05 USBν = 0.1

0.136±0.001 0.169±0.006 0.218±0.009

0.010±0.0005 0.006±0.002 0.008±0.003

0.144±0.004 0.206±0.010 0.261±0.012

0.005±0.0008 0.018±0.004 0.029±0.007

0.247±0.014 0.359±0.059 0.496±0.076

0.006±0.005 0.031±0.014 0.042±0.025

0.290±0.020 0.302±0.059 0.357±0.047

0.004±0.003 0.009±0.006 0.072±0.024

C.5. Gaussian Data We adopt the 1000D Gaussian Mixture Data from (Wang et al., 2025). Following their setup, at the initial time point, there are 500 cells from 2 Gaussians, 100 from the upper Gaussian, and 400 from the lower Gaussian. The upper Gaussian remains stationary and undergoes pure growth, whereas the lower Gaussian exhibits no growth and instead bifurcates symmetrically into left and right Gaussians. At the terminal time point, the total cell count is 1,400: the upper component increases from 100 to 1,000 cells, and the lower component splits evenly into two 200-cell Gaussians. Since the 1000D Gaussian has 1000 23

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

dimensions, it is suitable for testing USB’s performance in high dimensional spaces. As shown in Fig.5 and Fig.1, USB faithfully recovers the underlying dynamics. In the main text, we set δ = 1.4, ν = 0.001.

Figure 5. Learned trajectories on the Gaussian 1000D dataset (δ = 1.3, ν = 0.001).

C.6. Epithelial Mesenchymal Transition Data We adopt the dataset that captures the epithelial-mesenchymal transition (EMT) in A549 lung cancer cells from (Cook & Vanderhyden, 2020). The dataset includes samples collected at four distinct time points throughout this process. (Sha et al., 2024) reduced the dimension to 10 by an autoencoder. We applied USB on the reduced 10D EMT data. As shown in Table 7, USB achieves the best performance at most time points with growth penalty δ = 14, and diffusion parameter ν = 0.001. We plotted the learned growth rate in Figure 6. As shown, USB predicts higher growth rate in initial and intermediate stages, which is in line with the results predicted by DeepRUOT (Zhang et al., 2025) and WFR-FM (Peng et al., 2026) as these cells exhibit enhanced stemness. Table 7. Comparison of method performance over time on the 10D EMT dataset.

Method MMFM Metric FM SF2M MIOFlow BranchSBM TIGON DeepRUOT Var-RUOT UOT-FM VGFM WFR-FM USB

t=1

t=2

t=3

RME

0.2576 0.2605 0.2566 0.2439 0.3287 0.2433 0.2902 0.2540 0.2538 0.2350 0.2099 0.1831±0.0000

— — — — — 0.002 0.001 0.075 0.002 0.016 0.001 0.000±0.0000

0.2874 0.2971 0.2811 0.2665 0.3757 0.2661 0.3193 0.2670 0.2696 0.2420 0.2272 0.2159±0.0002

— — — — — 0.003 0.011 0.014 0.013 0.011 0.002 0.009±0.0000

0.3102 0.3050 0.2900 0.2841 0.2894 0.2847 0.3291 0.2683 0.2771 0.2450 0.2346 0.2309±0.0001

— — — — — 0.001 0.002 0.041 0.010 0.018 0.001 0.009±0.0000

C.7. Embryoid Bodies Data We adopt the EB dataset from (Moon et al., 2019), comprising 16,819 cells collected at five time points over a 27-day differentiation course of human embryoid bodies (EBs), an experimental model of early embryonic development. The data was preprocessed via Principal Component Analysis (PCA) as in (Wang et al., 2025). We kept the first 100 principal components for downstream analyses. Scalability with respect to the dimensionality. We evaluate USB on EB data with first 5, 50, 100 PCs to test the scalability 24

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots 1.0 0.90

0.8

Predicted rate

0.73

0.6

0.56

0.4

0.39

0.2

0.21

0.0

0.04

0.0

0.2

0.4

0.6

0.8

1.0

Figure 6. Learned growth rate on the EMT dataset (δ = 1.4, ν = 0.001).

of USB w.r.t the dimensionality. We set δ = 3, 27, 30 for 5D, 50D, 100D, respectively, and set ν = 0.001. To be consistent to (Peng et al., 2026), the 5D EB is further standardized. We compare the performance of USB with several other methods on the 3 datasets with results in showed Table 8,9,10, respectively. The computation time is recorded in Table 11. Across all experiments, USB achieved performance comparable to the best baseline with no significant computational overhead, indicating that USB scales directly to high-dimensional datasets without compromising performance. Table 8. Comparison of method performance over time on the 5D EB dataset. Best results are in bold, and the second best are underlined. Method MMFM Metric FM SF2M MIOFlow BranchSBM TIGON DeepRUOT Var-RUOT UOT-FM VGFM WFR-FM USB

t=1

t=2

t=3

t=4

RME

0.477 0.449 0.556 0.442 0.864 0.386 0.386 0.416 0.544 0.402 0.324 0.331±2 × 10−5

— — — — — 0.002 0.005 0.111 0.032 0.046 0.003 0.006±2 × 10−6

0.554 0.552 0.715 0.585 1.355 0.502 0.497 0.486 0.670 0.494 0.401 0.408±2 × 10−5

— — — — — 0.015 0.017 0.144 0.029 0.018 0.001 0.023±4 × 10−6

0.781 0.583 0.750 0.651 1.211 0.602 0.591 0.509 0.729 0.525 0.431 0.427±2 × 10−5

— — — — — 0.021 0.021 0.054 0.016 0.035 0.005 0.018±6 × 10−6

0.872 0.597 0.650 0.670 1.626 0.600 0.585 0.511 0.852 0.573 0.510 0.450±4 × 10−5

— — — — — 0.027 0.030 0.022 0.041 0.021 0.005 0.019±1 × 10−5

Table 9. Comparison of method performance over time on the 50D EB dataset. Best results are in bold, and the second best are underlined. Method MMFM Metric FM SF2M MIOFlow BranchSBM TIGON DeepRUOT Var-RUOT UOT-FM VGFM WFR-FM USB

t=1

t=2

t=3

t=4

RME

9.124 8.506 9.247 8.447 11.085 8.433 8.169 9.442 8.717 7.951 7.664 7.992±2 × 10−5

— — — — — 0.067 0.003 0.128 0.063 0.089 0.008 0.008±7 × 10−7

10.474 9.795 10.882 9.229 12.359 9.275 9.049 9.709 10.858 8.747 8.659 8.973±5 × 10−5

— — — — — 0.022 0.038 0.081 0.009 0.042 0.006 0.019±1 × 10−6

11.022 10.621 11.650 9.436 12.353 9.802 9.378 10.482 11.813 9.244 9.182 9.421±2 × 10−5

— — — — — 0.179 0.088 0.031 0.022 0.019 0.004 0.008±2 × 10−6

11.480 12.042 12.154 10.123 10.813 10.148 9.733 10.735 12.733 9.620 9.914 9.963±2 × 10−5

— — — — — 0.101 0.004 0.030 0.018 0.044 0.004 0.017±3 × 10−6

Sensitivity analysis for mini-batch OT. For larger datasets, full-batch OT can require substantial memory and long computation time. To avoid these costs, one can use mini-batch OT (Tong et al., 2024a; Fatras et al., 2021). To show that 25

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots Table 10. Comparison of method performance over time on the 100D EB dataset. Best results are in bold, and the second best are underlined. Method MMFM Metric FM SF2M MIOFlow BranchSBM TIGON DeepRUOT Var-RUOT UOT-FM VGFM WFR-FM USB

t=1

t=2

t=3

t=4

RME

11.460 10.806 11.333 11.387 12.853 10.547 10.256 11.746 10.757 10.313 9.941 10.206±5 × 10−5

— — — — — 0.014 0.002 0.091 0.056 0.048 0.009 0.016±5 × 10−7

13.879 12.348 12.982 12.331 14.298 12.926 11.103 12.237 12.799 11.278 11.040 11.297±6 × 10−5

— — — — — 0.052 0.074 0.024 0.037 0.035 0.006 0.010±8 × 10−7

14.441 13.622 13.718 11.905 14.419 13.897 11.529 12.957 13.761 11.703 11.516 11.758±9 × 10−6

— — — — — 0.107 0.136 0.150 0.044 0.028 0.008 0.002±2 × 10−6

14.907 16.801 14.945 12.908 13.718 14.945 12.406 13.335 15.657 12.637 12.664 12.888±5 × 10−5

— — — — — 0.096 0.047 0.074 0.022 0.066 0.005 0.003±3 × 10−6

Table 11. Computational time across different dimensions on the EB dataset.

Dimension

Time (s)

5 50 100

31.36 30.15 30.67

USB is compatible with this strategy—and thus has the potential to handle huge datasets—we evaluated different mini-batch sizes on the EB (100D) and Mouse (C.9) datasets. Results are summarized in Table 12 and Table 15. On EB (100D), we observe a non-monotonic dependence on batch size: using either small batches (500) or full-batch OT leads to slightly inferior performance and longer computation time, while a moderate batch size performs better and runs faster. Nevertheless, the overall impact of batch size is modest, and USB remains on par with the strongest baseline even with small batch size. This indicates that USB can naturally pairs with mini-batch OT and scales to huge datasets without losing performance. For consistency with the baseline using mini-batch OT (WFR-FM) and a favorable accuracy–efficiency balance, we use a mini-batch size of 2,000 on EB (100D), Cite, and Mouse datasets. Table 12. Sensitivity analysis for batch size of mini-batch WFR-OET on the 100D EB dataset. Batch Size

t=1

t=2

t=3

t=4

Time (s)

RME

500 1000 2000 3000 4000

10.238±3 × 10−5 10.228±2 × 10−5 10.206±5 × 10−5 10.201±4 × 10−5 10.189±1 × 10−5

0.015±4 × 10−7 0.014±8 × 10−7 0.016±5 × 10−7 0.015±6 × 10−7 0.016±9 × 10−7

11.360±3 × 10−5 11.314±6 × 10−5 11.297±6 × 10−5 11.337±3 × 10−5 11.306±3 × 10−5

0.010±4 × 10−7 0.008±2 × 10−6 0.010±8 × 10−7 0.010±1 × 10−6 0.011±1 × 10−6

11.877±5 × 10−5 11.789±6 × 10−5 11.758±9 × 10−6 11.803±4 × 10−5 11.756±3 × 10−5

0.001±1 × 10−6 0.005±6 × 10−5 0.002±2 × 10−6 0.004±1 × 10−6 0.002±6 × 10−7

12.958±6 × 10−5 12.880±2 × 10−4 12.888±5 × 10−5 12.905±5 × 10−5 12.816±4 × 10−5

0.006±2 × 10−6 0.001±3 × 10−6 0.003±3 × 10−6 0.001±2 × 10−6 0.003±1 × 10−6

30.33 27.76 27.26 28.34 28.42

w/o mini-batch

10.216±2 × 10−5

0.014±4 × 10−7

11.330±5 × 10−5

0.011±9 × 10−7

11.817±3 × 10−5

0.000±1 × 10−6

12.916±7 × 10−5

0.005±2 × 10−6

31.53

C.8. Cite-seq Data We adopt the CITE-seq (Cite) dataset from (Lance et al., 2022), consisting of 31,240 cells across four time points. Following the preprocessing in (Wang et al., 2025), we keep only the gene-expression modality and apply PCA to 50 dimensions. We use mini-batch OT (batch size 2,000) and set δ = 30, ν = 0.001.The performance on distribution matching and mass matching are summarized in Table 13. As shown, USB attains the highest accuracy in half of the experiments and remains comparable to the best baseline on the remainder. C.9. Mouse Hematopoiesis Data We adopt the mouse blood hematopoiesis dataset (Mouse) from (Weinreb et al., 2020), selecting 49,302 lineage-traced cells measured at three time points. The gene-expression matrix were reduced to 50 dimensions using PCA. Owing to its substantial population growth and large cell number, it is well suited for scalability testing and USB. We use mini-bath OT 26

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots Table 13. Comparison of method performance over time on the 50D CITE dataset. Best results are in bold, and the second best are underlined.

Method MMFM Metric FM SF2M MIOFlow BranchSBM TIGON DeepRUOT Var-RUOT UOT-FM VGFM WFR-FM USB

t=1

t=2

t=3

RME

33.971 28.314 29.543 28.290 39.908 28.196 28.245 30.219 33.531 29.449 27.831 27.021±5 × 10−5

— — — — — 0.186 0.168 0.331 0.009 0.020 0.043 0.002±1 × 10−6

36.854 28.617 32.655 28.524 37.008 27.921 27.908 32.702 32.795 29.722 27.478 28.081±1 × 10−4

— — — — — 0.545 0.525 0.325 0.046 0.057 0.045 0.015±2 × 10−6

43.721 33.212 36.265 32.230 31.883 32.846 32.950 40.613 49.751 33.752 34.784 38.447±0.003

— — — — — 0.653 0.634 0.486 0.097 0.001 0.022 0.010±5 × 10−5

of size 2000, and set δ = 15, ν = 0.001, the all time points training results is shown in Table 14. USB outperforms most baselines significantly except WFR-FM, and performs comparable to WFR-FM at all time points. Table 14. Comparison of method performance over time on the 50D Mouse dataset. Best results are in bold, and the second best are underlined.

Method MMFM Metric FM SF2M MIOFlow BranchSBM TIGON DeepRUOT Var-RUOT UOT-FM VGFM WFR-FM USB

t=1

t=2

RME

7.647 7.788 8.217 6.313 7.957 6.140 6.052 7.951 8.114 6.274 5.486 5.589±1 × 10−6

— — — — — 0.382 0.062 0.131 0.035 0.076 0.012 0.013±6 × 10−8

10.156 11.449 11.086 6.746 9.236 6.973 6.757 10.862 9.170 6.796 6.211 6.548±4 × 10−6

— — — — — 0.326 0.041 0.154 0.011 0.070 0.011 0.008±4 × 10−7

Sensitivity analysis for mini-batch OT. As mentioned in (C.7), we also test the sensitivity of the batch size of mini-batch OT on Mouse dataset. The result is shown in Table 15. The non-monotonic varying computation time is also observed. Intuitively, for small batch sizes, it take shorter time to compute each mini-batch OT, but the number of batch are larger, thus may also require more computation time. Hence, the computation time exhibits a U-shape variation according to batch size. On Mouse dataset, the W1 distance and RME decrease when using larger batch size while it slightly increase when using full-batch. The monotonicity is not observed on EB dataset. It may because of that EB is not large enough. Intuitively, the accuracy will be positively correlated to the batch size on sufficiently large datasets. Consistent to the results on EB, the performance and computation time show no significant difference across batch sizes which are not extremely small for Mouse (>1000). Scalability with respect to the cell number. We subsampled the Mouse dataset to 10000, 15000, 20000, 25000, 30000, 35000, 40000, 45000 cells to test the scalability of USB w.r.t the size of the dataset. The training time on each subsampled dataset are scattered in Fig.7. The training time scales linearly with cell numbers, in line with other competing methods, indicating that USB can be applied to larger datasets. 27

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots Table 15. Sensitivity analysis for batch size of mini-batch WFR-OET on the 50D Mouse dataset.

Batch Size

t=1

t=2

Time (s)

RME

500 1000 2000 3000 4000 5000 10000 20000

5.644±2 × 10−6 5.599±4 × 10−6 5.589±1 × 10−6 5.539±2 × 10−6 5.492±3 × 10−6 5.484±1 × 10−6 5.502±2 × 10−6 5.438±1 × 10−6

0.015±1 × 10−7 0.014±2 × 10−7 0.013±6 × 10−8 0.011±1 × 10−7 0.012±1 × 10−7 0.010±6 × 10−8 0.013±1 × 10−7 0.009±1 × 10−7

6.854±1 × 10−5 6.535±5 × 10−6 6.548±4 × 10−6 6.476±7 × 10−6 6.416±4 × 10−6 6.396±2 × 10−5 6.397±2 × 10−6 6.335±4 × 10−6

0.001±5 × 10−7 0.006±6 × 10−7 0.008±4 × 10−7 0.007±7 × 10−7 0.009±4 × 10−7 0.001±6 × 10−7 0.013±7 × 10−7 0.006±5 × 10−7

307.73 293.61 292.28 296.69 300.50 306.06 313.60 332.33

w/o mini-batch

5.527±8 × 10−7

0.008±1 × 10−7

6.383±1 × 10−5

0.001±5 × 10−7

331.22

Figure 7. Training time v.s. cell numbers

D. Relation to other algorithms D.1. Relation to SF2 M SF2 M (Tong et al., 2024b) is the first simulation-free framework for learning balanced Schrödinger Bridge between arbitrary source and target distributions. It utilized the KL-disintegration to decoupled the SB to two parts: the regularized OT (ROT) coupling induced by the static SB problem, and the conditional path bridging two Diracs. Following (Föllmer, 1988; Léonard, 2014), the static SB problem (7) can be written as a regularized OT (ROT) problem Z 1 KL(P01 ∥Q01 ) = 2 ∥x − y∥2 P01 (dx, dy) + KL(P01 ∥µ0 ⊗ µ1 ) + Cν (56) 2ν ⋆ Thus, SF2 M first calculated the optimal coupling π2ν 2 of the ROT problem Z 1 min 2 ∥x − y∥2 π(dx, dy) + KL(π∥µ0 ⊗ µ1 ) π 2ν

(57)

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

then obtained the SB P⋆ by integrating the conditional paths, which are Brownian bridges in this case Z ⋆ P⋆ = (νW)xy dπ2ν 2 (x, y)

(58)

SF2 M tried to find the most likely stochastic process referencing on a Brownian motion, while USB tries to find the most likely stochastic process referencing on a branching Brownian motion. In comparison, USB extends the ROT coupling to RUOT semi-coupling, and generalizes the Brownian bridge to Poisson-Brownian bridge to deal with unbalanced source and target. Though USB focuses on integrating the stochastic and unbalanced effect of single-cell dynamics, the balanced cases can be naturally included. Under balanced cases, by setting the reference branching rate λ → 0+ , USB reduces to SF2 M. To see that, one can calculate the no-growth limit of the growth penalty (46). Since the penalty is even, we assume g ≥ 0. p g g p lim Ψν,λ (g) = lim νλ(1 − 1 + g 2 /λ2 + ln( + 1 + g 2 /λ2 )) + + λ λ λ→0 λ→0 g g 2g −1 = lim+ νλ(1 − + o(λ ) + ln + o(λ−1 )) λ λ λ λ→0 (59) = lim (−νg + o(1) + νg ln(2g) − νg ln(λ)) λ→0+ ( 0, g = 0 = +∞, g ̸= 0 Hence, when λ → 0, the growth rate is forced to be 0, and the corresponding RUOT problem (10) reduces to Z 1Z 1 ∥u(x, t)∥22 ρt (x)dxdt RUOTno grwoth (µ0 , µ1 ) = inf ρ,g,u 0 X 2 ν2 ∆x ρ, ρ0 = µ0 , ρ1 = µ1 s.t. ∂t ρ + ∇x · (ρu) = 2

(60)

which is equivalent to the balanced SB problem (Chen et al., 2016; Pra, 1991). Hence the semi-coupling reduces to the coupling induced by the static SB problem. For conditional path, Poisson-Brownian bridge also reduces to the standard Brownian bridge when there is no growth. Thus, USB can reduce to SF2 M in mass conserved cases by referencing on a BBM with zero branching rate, which is a Brownian motion. D.2. Relation to WFR-FM WFR-FM (Peng et al., 2026) is an unbalanced extension of flow matching based on WFR geometry. It can learn both the velocity and the growth rate simultaneously without ODE simulation, and one can prove that the learned measure trajectory is a geodesic under WFR geometry. The algorithm efficiently addresses unbalanced mass in single-cell dynamics in both theory and practice, but it cannot model stochasticity. In contrast, USB extends WFR-FM’s unbalanced flow matching loss to a unbalanced score matching loss, thereby enabling efficient joint modeling of both unbalance and stochasticity in cell dynamics. Notably, USB cannot be simply regarded as a stochastic extension of WFR-FM; their microscopic dynamics are dissimilar. WFR-FM is based on WFR geometry, whereas USB is based on the BSB problem. The mass variation along WFR-FM’s conditional path is quadratic, causing the mass to decrease and then increase even when connecting two equal-mass Dirac measures via a WFR geodesic. By contrast, the mass variation along USB’s conditional path is exponential; when two equal-mass Dirac measures are connected by a Poisson–Brownian bridge, the mass does not change. Hence, USB cannot include WFR-FM as a limit case. One interesting question is that is there a stochastic version of WFR which could naturally reduce to WFR? D.3. Relation to DeepRUOT DeepRUOT (Zhang et al., 2025) employs dynamic RUOT to simultaneously model the stochasticity and unbalanced mass inherent in single-cell dynamics, solving the corresponding RUOT problem using a NeuralODE approach. In contrast, USB is proposed based on the BSB problem. It is also capable of simultaneously modeling stochasticity and unbalanced mass, and since it employs unbalanced score matching for solving, it is more efficient than DeepRUOT. To approximate the static semi-coupling within the BSB problem, USB also utilizes RUOT. However, it does not utilize DeepRUOT for the dynamic solution, instead, it approximates its static form using the quickly solvable static WFR. 29

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

D.4. Relation to UDSB UDSB (Pariset et al., 2023) is also an algorithm for solving stochastic transport between unbalanced measures. It is based on operator theory. UDSB introduces a coffin state outside the ambient space to model changes in total mass, and uses iterative proportional fitting (IPF) to solve the problem. In contrast, USB is designed based on the BSB problem, hence the two in fact tackle different problems. UDSB turns unbalanced stochastic transport into a diffusion Schrödinger bridge problem in an extended state space by introducing the coffin state; USB naturally introduces unbalanced effect by replacing the reference process of Schrödinger bridge from Brownian motion to branching Brownian motion. At the microscopic level, branching Brownian motion better matches the dynamical picture of cell division where new cells born from existing cells instead of a common coffin state, making USB more suitable for single-cell dynamics inference. Moreover, the solution techniques employed by the two methods differ substantially: UDSB is IPF-based and can be viewed as a unbalanced extension of classical Schrödinger bridge numerics, whereas USB is based on unbalanced score matching and should be regarded as a generalization of score matching methods. D.5. Relation to BranchSBM BranchSBM (Tang et al., 2025) is a recently proposed multimodal generalization of the SB problem, aimed at capturing branched or divergent evolution from a common origin to multiple distinct outcomes. It introduces mass weights into the generalized SB (Liu et al., 2024), thereby extending it to an unbalanced generalized SB; by training multiple such unbalanced generalized SB with conserved total mass across branches, BranchSBM models how a population splits from a single source into several branches and ultimately reaches multiple targets. The branching emphasized by BranchSBM is at population-level, focusing on the allocation and flow of mass among branches. Since the total mass is conserved, BranchSBM is inherently unable to model the unbalanced effect in single-cell dynamics. In contrast, USB emphasizes particle-level branching, reflecting the dynamical picture of cell division and apoptosis to model unbalanced effect. How to combine BranchSBM’s macroscopic branching structure with USB’s microscopic branching structure to build a framework that simultaneously models unbalance, stochasticity, and multimodality is an interesting direction for future research. D.6. Relation to VGFM VGFM (Wang et al., 2025) creatively employs semi-constrained OT and two-period transport to decouple the growth and velocity dynamics, subsequently achieving simultaneous modeling by integrating them through defined joint dynamics. In VGFM, the conditional path between Dirac pairs follows linear interpolation, while the mass evolution follows an exponential trajectory. This conditional path can be viewed as the limit of the USB conditional path as stochastic noise approaches zero; specifically, the reference process corresponds to a branching Brownian motion with a diffusion parameter of zero, exhibiting only birth-death without diffusion. However, VGFM cannot be regarded as a deterministic degeneration of USB due to differences in their adopted couplings. When calculating the coupling, VGFM utilizes a semi-constrained OT plan, which essentially performs OT between two balanced distributions, thereby treating transitions and mass variations separately at this stage. In contrast, USB is a principled algorithm based on the BSB problem. It employs RUOT semi-coupling to approximate the BSB problem, integrating transitions with birth-death dynamics within this very step. Consequently, even if the diffusion parameter of the reference process in USB is set to zero, it does not degenerate into VGFM. Furthermore, VGFM involves a two-stage training process: the first stage uses flow matching for warm-up, while the second requires optimizing an OT loss via NeuralODE, which is simulation-based, to achieve better performance. By comparison, USB is a fully simulation-free algorithm that requires only a single stage of unbalanced score matching training.

E. Discussion on general travelling dirac In this section, we first review Chizat’s method (Chizat et al., 2018a) for recasting WFR into the OET formulation—namely, by computing the WFR cost between two Dirac measures (the corresponding trajectory is called the travelling Dirac) and expressing the WFR cost between two distributions as an integral of the Dirac WFR cost, thereby converting dynamic WFR into static WFR and then obtaining OET via duality. We then discuss the difficulties this approach encounters in the general RUOT problem.

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

E.1. Travelling Dirac To solve a WFR problem i.e. a RUOT problem with quadratic growth penalty, one can reformulate the WFR problem as a integral of WFR costs between pair of Diracs. Z WFR2δ (µ0 , µ1 ) = inf WFR-DD2δ (γ0 (x, y)δx , γ1 (x, y)δy )dxdy, (61) 2 2 (γ0 ,γ1 )∈(M+ (X ))

where the WFR cost between two Diracs is m,x

s.t.

Z 1

1 ṁ(t) 2 (∥ẋ(t)∥22 + δ 2 | | )m(t)dt 2 m(t) m(0) = m0 , m(1) = m1 , x(0) = x0 , x(1) = x1

WFR-DD2δ (m0 δx0 , m1 δx1 ) = inf

(62)

Following (Chizat et al., 2018a), one can derive the Euler-Lagrange equation of (62). d ∂ d ∂  ( = )  (mẋ) = 0 dt ∂x dt ∂ ẋ 2 2 2  ∂ d ∂  1 |ẋ|2 − δ ṁ = δ 2 m̈m − ṁ ( = ) 2 2 2 2m m ∂m dt ∂ ṁ The first equation stands for momentum conservation, which yields

(63)

mẋ = ω

(64)

ω , the second equation becomes for some constant vector ω. By plugging in ẋ = m

2m̈m − ṁ2 = |ω|2 /δ 2

(65) 2

which describes the evolution of mass. Solutions of the second order ODE are of the form m(t) = At + Bt + C. Plugging in this ansatz, one can derive a set of equations of A, B, C.   C = m 0 A + B + C = m1 (66)   2 2 2 4AC − B = |ω| /δ The existence of the solution has also been proved by (Chizat et al., 2018a). This optimal transport trajectory between two Diracs is called travelling Dirac. Given the closed form of travelling Dirac, one can directly calculate (62) as √ ∥x0 − x1 ∥2 WFR-DD2δ (m0 δx0 , m1 δx1 ) = 2δ 2 (m0 + m1 − 2 m0 m1 cos( )) (67) 2δ Thank to this explicit form of Dirac-Dirac WFR cost, the WFR cost between two measures can be written in a static form Z p ∥x − y∥2 2 2 inf WFRδ (µ0 , µ1 ) = 2δ γ (x, y) + γ (x, y) − 2 γ (x, y)γ (x, y) cos( ) dxdy (68) 0 1 0 1 2δ (γ0 ,γ1 )∈(M+ (X 2 ))2 X 2 which is (4). Hence, the WFR problem (3) is decoupled into two parts: the static WFR semi-coupling and travelling Dirac. The former describes the mass transportation plan i.e. how much mass to be sent from source x and how much to be received at target y, while the latter describes what happens during the transportation process i.e. the variation of the position and mass of the specific Dirac given the source and target. The static WFR is further proved to be equivalent to the OET problem (50) which is easy to solve. E.2. The travelling Dirac for general RUOT For general RUOT problem (10), in order to obtain the semi-coupling, one need to find a static form for it. We tried to follow (Chizat et al., 2018a) to find the travelling Dirac of the general RUOT problem. We first write down the RUOT between two Diracs. Z 1 ṁ(t) 1 RUOT-DD(m0 δx0 , m1 δx1 ) = inf ∥ẋ(t)∥22 + Ψ( ) m(t)dt m,x 0 2 m(t) (69) s.t. m(0) = m0 , m(1) = m1 , x(0) = x0 , x(1) = x1 31

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

Then, we derive the Euler-Lagrange equation of it. d   (mẋ) = 0 dt 2  |ẋ|2 + Ψ − ṁ Ψ′ = m̈m − ṁ Ψ′′ m m2

∂ d ∂ = ) ∂x dt ∂ ẋ ∂ d ∂ ( = ) ∂m dt ∂ ṁ (

(70)

The first equation stands for the conservation of momentum, while the second equation again describes the evolution of ω mass. Plugging in ẋ = m , the second equation is of the form |ω|2 + m2 Ψ − mṁΨ′ = (m̈m − ṁ2 )Ψ′′

(71)

The complex second order ODE has no analytic solution for most of the choice of Ψ. If we want to derive an analytic solution, Ψ should be well chosen to make the above equation simple enough. To get rid of more complex terms, one may want Ψ′′ to be a constant instead of a function of ṁ m . Hence, for second order differentiable functions, a proper choice for Ψ is a polynomial with degree ≤ 2. Taking Ψ(z) = az 2 + bz + c, we have |ω|2 + cm2 = 2am̈m − aṁ2

(72)

When c = 0, it is WFR; when ac < 0, the mass evolves in cosine law; when ac > 0, the mass evolves in hyperbolic cosine law. Note that (72) is independent of the first order term coefficient b since the total cost induced by it is actually a constant R b 1 b 2 0 ṁ(t)dt = 2 (m1 − m0 ). For the RUOT problem with growth penalty like (46) or other general forms, the essential difficulty of deriving a static form for it is that the equation (72) is so hard to solve analytically. Though one can still write the RUOT cost between two measures as the integral of RUOT costs between Diracs Z RUOT(µ0 , µ1 ) = inf RUOT-DD(γ0 (x, y)δx , γ1 (x, y)δy )dxdy, (73) 2 2 (γ0 ,γ1 )∈(M+ (X ))

the RUOT cost between two Diracs (RUOT-DD) is intractable, hence it is challenging to solve the semi-coupling through this approach.

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

F. Algorithm workflow Algorithm 2 USB training workflow Require: Sample-able distributions µt0 , µt1 , . . . , µtK , diffusion parameter ν, branching rate λ, training batch size b, vector net vθ (x, t), growth rate net gθ (x, t), score net sθ (x, t). 1: for k = 0 → K − 1 do (k) (k) 2: (γ0 , γ1 ) ← coupling of RUOTν,λ (µk , µk+1 ) with penalty Ψν,λ (g) determined by (11) 3: end for 4: while Training do 5: for k = 0 → K − 1 do (k) 6: (xtk , xtk+1 ) ∼ γ0 (k) 7: t ∼ U(0, 1), t ← tk + (tk+1 − tk )t 8: ηt(k) ← xtk + t(x1 − xtk ) 9: x(k) ∼ N (ηt(k) , ν 2 t(1 − t)I) 1−2t 10: u(k) ← t(1−t) (x(k) − ηt(k) ) + (xtk+1 − xtk ) 11:

g (k) ← ln

12:

s(k) ←

mtk+1 (xtk ,xtk+1 ) mtk (xtk ,xtk+1 ) /(tk+1 − tk )

ηt(k) −x(k) ν 2 t(1−t) (k)

m(k) ← mt (xtk , xtk+1 )/γ0 (xtk , xtk+1 ) (Defined in 23 and 24) 14: end for c c c c c c 15: Concatenate {x(i) , t(i) , u(i) , g (i) , s(i) , m(i) }K i=1 into batch tensors {x , t , u , g , s , m } c c c 2 c c c 2 2 c c c 2 16: LCUSM (θ) ← (∥vθ (x , t ) − u ∥2 + ∥gθ (x , t ) − g ∥2 + λ (t) ∥sθ (x , t ) − s ∥2 )mc 17: θ ← Update(θ, ∇θ LCUSM (θ)) 18: end while 19: return vθ , gθ , and sθ 13:

Algorithm 3 Continuous inference workflow Require: Data at initial time point D0 , diffusion parameter ν, timestep ∆t, learned vector net vθ (x, t), growth rate net gθ (x, t), score net sθ (x, t). 1: for x in D0 do 2: ω=1 3: while simulation do 4: ξ ∼ N (0,∆tI) 2

x ← x + vθ (x, t) + ν2 sθ (x, t) ∆t + νξ

ω ← ω · egθ (x,t)∆t end while 8: Dweight append (x, ω) 9: end for 10: return Dweight 6:

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

Algorithm 4 Branching inference workflow Require: Data at initial time point D0 , diffusion parameter ν, timestep ∆t, learned vector net vθ (x, t), growth rate net gθ (x, t), score net sθ (x, t). 1: Dnext ← D0 2: while simulation do 3: D ← Dnext 4: for x in D do 5: ξ ∼ N (0,∆tI) 2

x ← x + vθ (x, t) + ν2 sθ (x, t) ∆t + νξ

α ∼ U [0, 1] if α ≤ |gθ (x, t)∆t| then 9: if gθ (x, t) ≥ 0 then 10: Dnext append x 11: Dnext append x (x divides into two) 12: else 13: pass (x dies) 14: end if 15: else 16: Dnext append x (No branching) 17: end if 18: end for 19: end while 20: for x in Dnext do 21: Dweight append (x, 1) 22: end for 23: return Dweight 7: 8:

G. The Use of Large Language Models (LLMs) LLMs were used only for grammatical correction to enhance the overall readability of this paper.

Record · ID 151756 · SHA-256 f1376a1e11fb21d2
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.