Modeling Cell Dynamics and Interactions with Unbalanced Mean Field Schrödinger Bridge
Authors: Zhenyi Zhang
Zihan Wang
Yuhao Sun
Tiejun Li
Peijie Zhou
Venue: NeurIPS 2025 poster ()
Keywords: optimal transport
Schrödinger bridge
trajectory inference
single-cell
Primary area: machine_learning_for_sciences
TL;DR: We introduce CytoBridge, a new deep learning approach that models unbalanced stochastic cellular dynamics and interactions from snapshot data.

Abstract: Modeling the dynamics from sparsely time-resolved snapshot data is crucial for understanding complex cellular processes and behavior. Existing methods leverage optimal transport, Schrödinger bridge theory, or their variants to simultaneously infer stochastic, unbalanced dynamics from snapshot data. However, these approaches remain limited in their ability to account for cell-cell interactions. This integration is essential in real-world scenarios since intercellular communications are fundamental life processes and can influence cell state-transition dynamics. To address this challenge, we formulate the Unbalanced Mean-Field Schrödinger Bridge (UMFSB) framework to model unbalanced stochastic interaction dynamics from snapshot data. Inspired by this framework, we further propose CytoBridge, a deep learning algorithm designed to approximate the UMFSB problem. By explicitly modeling cellular transitions, proliferation, and interactions through neural networks, CytoBridge offers the flexibility to learn these processes directly from data. The effectiveness of our method has been extensively validated using both synthetic gene regulatory data and real scRNA-seq datasets. Compared to existing methods, CytoBridge identifies growth, transition, and interaction patterns, eliminates false transitions, and reconstructs the developmental landscape with greater accuracy. Code is available at: https://github.com/zhenyiizhang/CytoBridge-NeurIPS.

Modeling Cell Dynamics and Interactions with Unbalanced Mean Field Schrödinger Bridge
Zhenyi Zhang Zihan Wang"." Yuhao Sun Tiejun Li and Peijie Zhou 2.3,4,..
LMAM and School of Mathematical Sciences, Peking University.
2 Center for Quantitative Biology, Peking University.
Center for Machine Learning Research, Peking University.
*EELDA, Peking University. AI for Science Institute, Beijing,
Emails: tzhenyizhang, jackwzh, 250111 1111111124)))) @stu. pku. edu. cn, {tieli, , (anjjzooo@ppuut edu.
Abstract
Modeling the dynamics from sparsely time-resolved snapshot data crucial for un- derstanding complex cellular processes and behavior. Existing methods leverage op- timal transport, Schrödinger bridge theory, or their variants to simultaneously infer stochastic, unbalanced dynamics from snapshot data. However, these approaches re- main limited in their ability to account for cell-cell interactions. This integration is essential in real-world scenarios since intercellular communications are fundamen- life processes and can influence cell state- transttranittion dynamics. To address this challenge, we formulate the Unbalanced Mean-Fiee Field Schrödinger Bridge (UMFSB) framework model unbalanced stochastic interaction dynamics from snapshot data. Inspired by this framework, we further propose CytoBridge, a deep learning algorithm designed to approximate the UMFSB problem. By explicitly model- ing cellular transitions, proliferation, and interactions through neural networks, CytoBridge offers the flexibility to learn these processes directly from data. The fectiveness of our method has been extensively validated using both synthetic gene regulatory data and real ScRNA-seq datasets. Compared to existing methods, Cyto- Bridge identifies growth, transition, and interaction patterss, eliminates false transi- tions, and reconstructs the developmental landscape with greater accuracy. Code is available at: https: //gpt/h/ com/ //ee/ieeiiii izhang/ cechannyyytorrdddeec CytoBridge NeurIPS. -
1 Introduction
Reconstructing dynamics from high- nigh-dimensional distribution samples is central challenge in science and machine learning. In generative models, methods such as Variational Autoencoders (VAEs) (Kingma and Welling 2013), diffusion models (Ho et al. 2020: Sohl-Dickstein et al. 2015; Song et al. 2021), and flow matching (Lipman et 2023; Tong et al. 2024a) have achieved success in generating high-fidelity images by coupling nigh-dimensional distributions (Liu et al. 2023). Meanwhile, in biology, inferring dynamics (also known as trajectory inference problem) from several static snapshots of single-cell RNA sequencing (scRNA-seq) data (Ding et al. 2022) to build the continuous dynamics of an individual cell and construct the corresponding cell-fate landscapes has also attracted broad interests (Schiebinger et al. 2019; Klein et al. 2025; Zhang et al. 2025b).
To study the trajectory inference problem, optimal transport (OT) theory serves as foundational tool (Bunne et al. 2024; Zhang et al. 2025b; Heitz et al. 2024; Zhang et al. 2025c). In particular, several works propose to infer continuous cellular dynamics over time by employing the Benamou-Brent
'These authors contributed equally. 'Coresponding authors.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

formulation (Benamou and Brenier 2000). Given that cell growth and death are critical biological processes, modeling the coupling of underlying unnormalized distributions has led to the development of unbalanced dynamical OT by introducing Chizat et al. 2018b). To further account for the prevalent stochastic effects on single-cell level, methods inspired by the Schrödinger Bridge (SB) problem seek to identify the most likely stochastic transition path between two arbitrary distributions (Gentil et al. 2017; Léonard 2014). To tackle both stochastic and unbalanced effects simultaneously, methods have been developed to model unbalanced stochastic dynamics (Pariset et al. 2023; Lavenant et al. 2024), along with a recent deep learning method (Zhang et al. 2025a), which leverages the regularized unbalanced optimal transport (RUOT) framework to infer continuous unbalanced stochastic dynamics from samples without requiring prior knowledge.
Masserstein-Fisher--ar
metric (Chizat et al. 201 8a;
Nevertheless, the majority of existing trajectory inference methods do not account for cell-cell interactions in cell-state transition dynamics, which involve important biological processes such as intercellular communications (Almet al. 2024; Tejada-Lapueed et al. 2025; Cang et al. 2023). Developing frameworks to infer unbalanced and stochastic continuous dynamics with particle interac- tions from multiple snapshot data remains a critical yet underexplored challenge.
To address this challenge, we propose the Unbalanced Mean Field Schrödinger Bridge (UMFSB), a modeling framework based on the Mean Field Schrödinger Bridge that extends to unnormalized distributions. We further develop new deep learning method (CytoBridge) to approximate the general UMFSB and learn continuous stochastic dynamics with cellular interactions from snapshot data with unbalanced distributions. Our primary contributions are summarized as follows:
We formulate the UMFSB problem to model unbalanced stochastic dynamics of interactive particles from snapshot data. By reformulating UMFSB with Fisher regularization form, we transform the original stochastic differential equation (SDE) constraints into computationally more tractable ordinary differential equation (ODE) constraints.
We propose CytoBridge, a deep learning algorithm to approximate the UMFSB problem. By explicitly modeling cellular growth/death and interaction terms via neural networks, CytoBridge does not need prior knowledge of these functions.
We validate the effectiveness of CytoBridge extensively on synthetic and real scRNA-se datasets, demonstrating promising performance over existing trajectory inference methods.
2 Related Works
Various Dynamical OT Extensions and Deep Learning Solvers Numerous efforts have been made to learn dynamics from snapshot data. To tackle the dynamical optimal transport problem, several methods have been proposed by leveraging neural ODEs or flow matching techniques (Tong et al. 2020; Huguet et al. 2022: Wan et al. 2023; Zhang et al. 2024a; Tong et al. 2024a; Albergo et al. 2023; Palma et al. 2025; Rohbeck et al. 2025; Petrović et al. 2025). To account for sink and source terms in unnormalized distributions, Peng et al. 2024; Sha et al. 2024: Tong et al. 2023; Eyring et al. 2024 developed the neural ODE-based solver for unbalanced dynamical OT. Wang et al. 2025 developed a flow matching approach to simultaneously learn velocity and growth. For the Schrödinger Bridge (SB) problem, approaches have been proposed based on its static or dynamic formulations (Shi et al. 2024; De Bortoli et al. 2021; Gu et al. 2025; Koshizuka and Sato 2023; Neklyudov et al. 2023; Neklyudov et al. 2024: Zhang et al. 2024b; Bunne et al. 2023; Chen al. 2022b; Zhou et al. 2024a; Zhu et al. 2024; Yeo et al. 2021; Jiang and Wan 2024), with corresponding flow matching methods (Tong et al. 2024b). To further incorporate unbalanced effects in the SB framework, methods utilizing branching SDE theory (Lavenant et al. 2024; Ventre et al. 2023; Chizat al. 2022),
backward SDE (Pariset al. 2023), neural ODEs with Fisher information regularization (Zhang et al. 2025a), or with first-order optimality conditions (Sun et al. 2025) have been introduced. However, theoretical formulation, along with an effective deep learning solver to simultaneously account for unbalanced stochastic effects and particle interactions in the dynamical OT framework, remains largely lacking.
Modeling Cellular Interactions in Trajectory Inference Several studies have explored the incor- For instance,
poration of cellular interaction effects into time-series ScRNA-seq trajectory inference. Atanackovic et al. 2025 employs a graph convolutional network within flow- now-matching framework

---

to model the impacts of neighborhood cells within the initial cell population. You et al. 2024 intro- duced a population-Ievel regularization in the energy form. Yang 2025 formulated the topological Schrödinger Bridge problem on a discrete graph. Fu et al. 2025 improves the accuracy of pseudotime inference by integrating cellular communication pattern. However, an explicit quantification of cell interaction dynamics in scRNA-seg trajectory inference is yet to be explored.
Mean-Field Control Problem Several works have explored mean-field problem (Zhou et al. 2024b; Ruthotto et al. 2020; Lu et al. 2024; Yang et al. 2024; Huang et al. 2024; Han et al. 2024; Shen and Wang 2023; Li and Liu 2025; Li et al. 2023; Shen et al. 2022) and its variants, as well as the incorporation of particle interaction terms in the Schrödinger Bridge Problem. Backhoff al. 2020; Hernández and Tangpi 2025 investigated theoretical properties of the Mean Field Schrödinger Bridge Problem. Rapakoulias al. 2025 develop deep learning solver for the MFSB problem. Liu et al. 2022b; Liu et al. 2024 proposed the generalized formulation of Schrödinger bridges that includes interacting terms. Yang et al. 2022 formulated the ensemble regression problem and developed neural ODE-based approach to learn the dynamics of interacting particle systems from distribution data. However, these approaches either require prior knowledge to specify the interaction potential field or do not account for unbalanced stochastic dynamics.
3 Preliminaries and Backgrounds
In this section, we provide an overview unbalanced stochastic effects and interaction forms within the dynamical OT framework. By integrating two perspectives of RUOT and MFSB described below, we motivate the formulation of the Unbalanced Mean Field Schrödinger Bridge (UMFSB) framework.
3.1 Regularized Unbalanced Optimal Transport
The regularized unbalanced optimal transport, also known as the unbalanced Schrödinger Bridge problem (Chen et 2022c), considers both the unbalanced stochastic effects in the dynamical OT framework (Baradat and Lavenant 2021; Zhang et al. 2025a):
Definition 3.1 (Regularized Unbalanced Optimal Transport). Consider
inf ('" (p.b.g) Jge 2
T
/
1
( 1122 p(x, ()dxdl ['' J av (g(x,()) p(xx l)dxdl,
where V R [0, +o0) corresponds to the growth penalty function, and a the weight of the growth penalty. The infimum is taken over all pairs (p, b. 9) such that P)-:0) = Vo, pl-, 1) = 2/1 p(x, t) absolutely continuous, and
Opp= -V (pb) - 1 (t) + gp
2
with vanishing boundary condition:
lim p(x,t) = 0.
[x] 1x00
Here bxx, t) is the velocity. g(x, f) is the growth function, and o (t) is the diffusion rate. Note that here ii and in1 are not necessarily the normalized probability densities, but are generally unnormalized densities of masses.
3.2 Mean Field Schrödinger Bridge Problem
Schrödinger bridge problem aims to find the most probable path between a given initial distribution No and target distribution V1, relative to reference process. Formally, it can be stated as:
min
DKL
(Mo.nyl4%..1)
Mo Po. A
=P
X where 1 is the probability measure induced by X, (0 st 1) and the reference measure Pf..1] However, the classical Schrödinger bridge problem considers the independent particles. The mean field extends the SB problem to the interacting particles with given initial and final distributions. We consider the dY, o t),,, where o (Y,,t) E dddd is the diffusion rate, W, I2 is
Y
given
3

---

Brownian motion and it is called the diffusion Schrödinger bridge problem, where it has a dynamic formulation. So the mean field Schrödinger bridge problem can be stated through this dynamical formulation (Backhoff et al. 2020; Hernández and Tangpi 2025):
the standard
Definition 3.2 (Mean Field Schrödinger Bridge Problem). Consider
lb 1 ()x(x l)dldx, inf (b.p)
(I)
The infinium is taken over all (b, p) subject to p(x; 0) vu. p(xx 1) vn, and
ip(x,t) - [(b(x.))) - / Rd 0t
k(x, V f(x y)ply, 1) dy) p(x, t) + o²(t)
Ap(x, t),
2
(2)
where & is the interaction potential and it is satisfied (oeeedd/((30 ((-x)=(((> The k(,,)) IRd x TRd + R is the interaction weight function.
4 Unbalanced Mean Field Schrödinger Bridge
In this section, we introduce the unbalanced mean field Schrödinger Bridge problem. Inspired by regularized unbalanced optimal transport, the dynamical formulation Definition 3.2 suggests a natural way to relax the mass constraint by introducing growth/death term go in (2). Meanwhile, we also introduce a loss function in (1) which considers both the growth and transition metric.
Definition 4.1 (Unbalanced Mean Field Schrödinger Bridge). Consider
inf / (b.9.P:0) Rd ///"" | 1)xx dddddx + Sea.h 1'" V(g(xxt))xxx l)dldx,
T
where (f():: IR + R+ is the growth cost function. The infinium is taken over all (b, p, 4) subject to p(x, 0) = Vu. p(x, 1) = 11, and
Op(x,t) t)
(D(x.0) - - (x))sxxxxx))x,,.. t) p(x, 1) o2(1) Ap(x, t)+gp,
ØL = Vx. -
where & is the interaction potential and is satisified o(-x)= - -xx).. The k(.))) : Rdd Rd R is the interaction weight function.
In Definition 4.1, if kxxx, y) = O. which means there is no cell-cell interaction, degenerates the regularized unbalanced optimal transport problem. If the growth penalty is set such that 9 t) must be zero (i.e., by setting Y (g(x. t)) +o for g(x, t) 0 and (0) = 0). the framework degenerates to simpler forms: if interactions are present (k(x, * 0), the formulation reduces to the Mean Field Schrödinger Bridge problem; if interactions are absent (ki(x, = O), reduces to the regularized optimal transport problem. It becomes the unbalanced dynamics optimal transport when k(x,y ) = (t) = and (g)) = g² becomes the dynamics optimal transport when growth, interaction and diffusion all goes to zero. We can reformulate Definition 4.1 with the following Fisher information regularization.
Theorem 4.1. The unbalanced mean field Schrödinger Bridge Definition 4. equivalent to
inf 1-' Jel 1il2 + x²(t) | lox (p.b.0,)))
pll + 2 ov),0 log + all (g) p(x, t)dxdl,
(3)
where V(): - Th is the growth cost cost function, and a is the weight of the growth cost. The infinium is taken over all (b.g.P b) subject p(x, 0) vo, p(x, V1, and
0plx,t) Vx (v(x. - Øt
v)x y)p(y, dy p(xxt) + +gp. -f
(4)
where t is the interaction potential and is satisified क6 ((-x) xx)). The k(.,)) )) TRd x Rdd the interaction weight function.

---

Here v(x, t) is new vector field function. This Fisher information form transforms the original SDE problem into the ODE problem which is computationally more tractable. Next, we will focus on solving this problem. The proof is simple and we left it in Appendix D.1 for reference.
Remark 4.1. From the proof of Theorem 4.1, the relation between the new vector filed v(x. t) and b(x, 1) is v(x. t) = b(x, t) 2 =² (t) log p(x, t). The new v(x. t) is also known as probability flow ODE and V log p(x.,) is the score function. Conversely, if the probability flow ODE v(x, t) and the score function Vx log p(x, 1) are known, then we can recover the original drift term.
Learning Cell Dynamies and Interactions through Neural Networks
Assume that we collect scRNA-seq samples from unnormalized distributions X, E Ru.xd (t 1.2. 7) at T time points, where n; is the number cells at time i and d is the number of genes, here we propose the CytoBridge algorithm to approximate UMFSB through neural networks. We parameterize transition velocity v(x, f), cell growth rate g(x, t), log density function 2 (t) log p(x, t) and cellular interaction potential $(x, t) using neural networks Vg. 9s. 5g and & respectively, as shown in Fig. 1. To effectively approximate the loss function in Theo- rem 4.1, we model the evolution of the mass densities through number of weighted interacting particles, which is supported by the following proposition.
Time
Time death
birth
noise
deda!!
[[vx0 1)-
velocity
Score
ODE solver
Proposition 5.1. Consider a system of weighted particles in Rd, where each particle has position X; E Rd and a positive weight w.(() > 0. The weight wi(l) evolves according to the ordinary differ- dus de
Interaction
growth
Figure 1: Overview of CytoBridge.
ential equation (ODE)
= g(Xi,,) t)),())),,,))....)) where
g I& [] + I is a given growth rate function. stochastic differential equation (SDE)
The position X evolves according to the
1
x = b(X;,
K(X) X)), X dt +o()
N N
3#1
Under assumptions stated in D.I, in the limit of N - co, the weighted empirical measure ju N EN w; w,(()xx; converges weakly to deterministic measure p(x, 1) dx, where p(x, 1) a weak solution to the partial differential equation (PDE)
Op
Øt (00=(((bx..)) [ x)
-
1x y)o(y, 1) dy p(x pbx.0)
o²(t) o2(0) 2
with the initial condition p(x, 0) (x()).
The derivation is left in Appendix D.2. By combining the results in Proposition 5.1 and Theorem 4.1, the ODE constraints we simulate is indeed dx; v(Xi, t)dt - Σ E ki, VxD(Xi
X dt. where ki k(X,, X is the cell-cell interaction strength and w, w (X(:) is the particle weight.
5.1 Simulating ODEs: Random Batch Methods
In ODE simulation, the computation complexity is O(22) due to the cellular interaction term. To speed up simulation, we adopt the Random Batch Methods (RBMs) (Jin et al. 2020; Jin et al. 2021; Jin and Li 2022), which transforms the interaction among all particles into interaction with particles within random grouping only. The algorithm reduces the computational complexity to O(pp I). where is the number in the batch C,, Assuming v(x, f) and & satisfy certain conditions, it has the
(1) P (1), PN (1) < CVT (Jin et al. 2020; Jin et al. 2021), where (1) (L) the empirical distribution produced by Algorithm and n (1) the distribution by simulating the original ODEs.
convergence result such that W3
the step size. The /
5

---

Algorithm Simulating ODEs: RBM
1: for m 1 [T/T] do
Randomly divide (1,2,-. .N = np} into n batches
2: for each batch do
N Σ jeIJEC,
Update particles with dx; = v(Xi, tjdt
x/)dt
5.2 Reformulating the Loss with Weighted Particle Simulation
The total loss to solve Theorem 4.1 is composed of three parts, e-g., the energy loss, the reconstruction loss, and the Fokker-Planck constraint such that & = Linergy + L + Xy Here Faergy loss promotes the least action principle of transition energy, CRecons promotes the matching loss. i.e., p(x,1) 1) V1, and LFP promotes the three neural networks that satisfy the Fokker-Plance equation constraint. We reformulate the loss terms through weighted particle representation and an RB.M-based Neural ODE solver.
Energy Loss Generalizing the idea in (Sha al. 2024), the energy loss in CytoBridge is equivalent Exarpo J0 2 llvo(x(t), t) + 2 V v (vo(x(t), t), so) a (oo) wo (t)dd,
where we (t)
exp Jo Jo go (x(t),)) and x(t) satisfies dx /dt
v(xit)
N-1 Ejtr k(xxxxx))x) V Vx x'). However, direct optimization of this term is challenging due to the involvement of the inner product (v)(x()), 1), 8a) which introduces mutual dependencies between the optimization of v and sg. To address this issue and simplify the computation, adopt an upper bound of the energy for training purposes (Appendix A.5).
Reconstruction Loss The reconstruction loss aims to align the final generated density to the true data density. Here we need to consider the unbalanced effect, so we use the unbalanced optimal transport align it. LResons Xm.-Mase + AaLor- The LMas is used to obtain the cell weights and align the cell number/ mass in the datasets. We then use the weights to normalize the distribution and apply Lor to match the distribution. In this work, we employ the local mass matching strategy from (Zhang et al. 2025a). Specifically, the trajectory mapping function ob predicts particle coordinates A Ar-1 from initial set Au over time indices T. governed by dx dx/dl = Va if no interaction is considered, or with a modified velocity incorporating the interaction potential Φ: v(x², () = 1 N- Ejzt ki(x', V 0 x²). Additionally, weight mapping function oy models particle weights via log w. (t)/tdt = (t,, t). starting from initial weights (0) I/N where N represents batch size. Mathematically, we use the empirical measure 2 = + Ei= dz; to approximate the true distribution po. hence the uniform weights. This convergence to the true distribution is guaranteed when N + xo. The mass matching loss Laass is composed of two terms. The first term is defined as the local mass matching loss Liocal Mass = Ex= MR. where Mr quantifies the error between predicted weights w;(4x) and target weights based on the cardinality of mapped points. As detailed in Appendix A.4, the target weights, derived from the number of closest real data points, encourages the growth network to assign higher weights to particles moving into denser state space regions, thus provides fine- fine-grained guidance on the growth network. Besides, the second term is defined as the global mass matching Ei= Gk. where GR is used to align the change of weights in total. An optimal transport loss for Ek W2(w", w(fx)) further aligns the predicted and observed distributions. Details can be found in Appendix A.4.
loss LGiohal
Fokker- Planck Constraint To enforce the physical relationships among the four neural networks, it is essential to introduce a physics-informe loss (PINN-loss), which incorporates the Bokker-Plancke constraint as guiding principle. LEP = |/ppo + Vx (poVo) - -oppll + A Ix 0) poll. where 2 Po cxp 72 sg.
Training CytoBridge aims to train four neural networks to model cell dynamics and interactions. To stabilize the training procedure, we leverage a two-phase training strategy. In the pre-training

---

phase, we seek to provide suitable initialization of these four neural networks. We first initialize Go and Vo without the interaction term to provide an upproximated matching. Then we train D, with fixed yo, leading to refined dynamics. The score network 8p is trained based on conditional flow matching. the training phase, these initialized networks are further refined by minimizing the proposed total loss. We summarize the training procedure in Algorithm 2 and Appendix A. We conduct ablation studies on different components of our training procedure in Appendix B.7. The selection of loss weighting is discussed in Appendix C.2.
Algorithm 2 Training CytoBridge
Require: Datasets Ag, ... Ar-1. batch size N, ODE iteration Rede. log density iteration rloogdensity
Ensure: Trained neural ODE Ve, growth function ge, score network Se and interaction Dg.
1: Pre-Training Phase:
2: for = 1 to nede do
b 1. Initialize growth go and velocity vo
3:
for t= to T - do
4:
Ac+1- t 0% (Ar t 1), w/ ((A+1)) + [g)(
5:
CRecOns + Me + WWW2(W", :(()); update vo. g0 w.r.t. CRecos
6:
for i = 1 to Rock do
>> 2. Initialize interaction potential D
7:
for fr T do
8:
Ac+1 +- (A,, + 1), w/ w(A++1) + dg((v(A)),
9:
CRocONs t -W2(w* w(t);; update va, Pg w.r.t. LRecons
10: for (= t T 2 do A t (A,,, 1)
D. Generate datasets Å.. D 4. Initialize the score network so
11: for i = to Mog demsity do
x1) (X., x1), u(0 1). x~ p(x,t t No, x1) using Ao. 13: Cscore + A V xx, t) + E 11:: update w.r.t. Lzcore
14: Training Phase:
Estimate initial distribution (x) from Ag using Gaussian Mixture Model (GMM).
16:
for = 1 to nldde do
for f- - do
18:
Aet1 +- (Ae, l + 1), w( ((A+1)) +- (g(
Linergy + -x, r== 2 1 //ll1/ Vx50112 + ||olll||solla + an (ow) wo(
19:
20:
LRecons +- +- Am + G1) + AdWa(w' w(t)
21:
LFP + 110,p((x,T)
Vx
T))
T)
Aw ||pp(x, 0) Po (x)
22:
LTaMl &- Linergy + Xx PReemns + LFP: update vo, 9o. 80, $o w.r.t Lratl
6
Results
Next, we evaluate CytoBridge's ability to simultaneously learn cell dynamics and cell-cell interactions. In computations below, we take (gg(x, t)) g (x,t) and (t) is constant. We also assume the interaction term is dependent on the distance between cells in gene expression space and we use radial basis functions (RBFs) to approximate it (Appendix A.3).
Synthetic Gene Regulatory Network In order examine CytoBridge's capabilities of learning cell dynamics as well as their underlying interactions simultaneously. we conducted experiments on the three-gene simulation model following (Zhang et al. 2025a). The dynamics of the original three- gene model are governed by stochastic ordinary differential equations, incorporating self- actvvacti,,ion. mutual inhibition, and external activation (Appendix B) as shown in Fig. (a). Additionally, we incorporated interactions into the simulation process. We consider the following types of interactions: potential (both attractive and repulsive) (3) no interactions. We aim to test whether CytoBridge can recover interaction in each case. For attractive interactions, the cells with similar gene expressions tend to converge toward similar levels. The interaction potential + is defined as: +(x - y) ||x yl As shown Fig. 2(b), the dynamics of the three-gene model exhibit quiescent area well as an area with notable transition and
(I) attractive interactions. (2) Lennard-Jones

---

increasing cell numbers. Moreover, the attractive potential results in the reduction in variance of the observed data at different time points. We compared CytoBridge with other methods across all time points using the Wasserstein distance (W1) and the Total Mass Variation (TMV) metric, defined in Appendix C. We summarized the results Table 1. is shown that CytoBridge achieves the best performance in both distribution matching and mass matching. Balanced Schrödinger bridge (e.g.. (Tong et al. 2024b)), which neglects both growth and interaction terms, results in false transition and variance patterns (Fig. 2(c)). DeepRUOT (Zhang et al. 2025a), which is an unbalanced Schrödinger Bridge solver, exhibits correct transition patterns but fails to capture the reduction in variance as it neglects the cell-cell interactions (Fig. 2(d)). By leveraging the UMFSB framework, shown in Fig. 2(e), CytoBridge correctly models both the transition patterns and the reduction in variance as the correct interaction potential (Fig. 2(f)) and growth rate (Fig. 2(g)) can be directly learned by CytoBridge. Also, CytoBridge is capable of constructing the underlying Waddington landscape (Fig. 2(h)). The low-lying regions this landscape correspond areas of high probability density, representing stable states. Other than the attractive interaction potential, we also incorporated the potential, and interaction case. Results can be found in Appendix 1.1, Figs. 4 and 5 and we find CytoBridge can correctly identify both the LJ potential and no interaction in these cases. Overall, both quantitative results and qualitative analysis indicate the necessity of incorporating both growth and interaction terms.
(e)
Transition Growth
(b)
True dynamics
(c)
Balanced SB
(d) DeepRUOT (no interaction)
X1
X2
Gume X,
Gene Xs
cene NG
X3
Interaction
Transition X
Transition V
140
Variance
x
Viariane
Quiescent
Gene X,
Gene X3
(o)
Ours
(9)
Learnt interaction
(9)
Learnt growth rate
(h)
Learnt landscape
Gene Xg
Gene
Transition
Veriance
OUT1
0.0
Gene
Ge
Figure 2: (a) Illustration of the synthetic gene regulatory dynamics. (b) The ground truth cellular dynamics project on (X1, X2). The red lines indicate the ground truth trajectories of cells in (b), or inferred trajectories of cells in (c) to (c). (c) The dynamics learned by balanced Schrödinger bridge SF2M (Tong et al. 2024b). (d) The dynamics learned by DeepRUOT. (e) The dynamics learned by CytoBridge. (f) The learned interaction potential.
(g) The growth rates inferred by CytoBridge. (h) The constructed landscape & 4. The z-axis represents the density of cells.
Table 1:
Wasserstein distance (Wi) and Total Mass Variation (TMV) of predictions at different time points across five runs synthetic gene regulatory data with attractive interactions (o 0.05). We show the mean value with one standard deviation, where bold indicates the best among all algorithms.
2
t
4
Model
W/
TMV
Wi;
TMV
F1:
TMV
W:
TMV
SP2M (Tong et al. 2024b)
0.146_ax0
0.080_axe
a. 3M
0,250_ame
0.447
0
0.584gmes
0.930gmmm
Meta FM (Atanackovic et al. 2025)
0.149=aax
0.080-1xx
0.241=axc
0.250
0.288
0,515=2000
0.404=aum3
0.930zu0o
MMFM (Rahbeck et 2025)
0.101
0.080_
0.223_axe
0.250_ave
0.4
0.5152m00
0.366ga0e
2.930
Metric FM (Kapusnink et 2024)
0,319:am<
0.080-axc
0.751-axx
0.250-am
0.690
0.515
0.930 0.614:suz
UOT-FM (Eyring al. 2024)
0,051_ax<
0.010_1xx
0.058_axx
0.036_axx
0.060
0.044_aam
0.054gm0ma
2.095
M0OFlew (Haguet et al. 2022)
0.315-axx
0.080-axx
0.387
0,250 aax
0.483 -000
0.515-aur
0.518
0.930-c00
uAM (Neklyadov et al. 2023)
0.489=axx
0,081=axx
0.995=x3x
0.033
1,402
0.459
1.655
I 1
UDSB (Pariset et 2023)
1.131_ax
0.
1,489_m
0.135_am1
1,455_0m
0.447
0.543
1.01 20nos
TIGON (Sha al. 2024)
0.169
0.097=a3c
0.184=aac
0.165-amm
0.167-aaxd
0.210
0.179esuma
0. 0.384:3000
DST (Zhang 2025a)
0.044_333
0.014_
0.045_130
0.026
g.053_aaee
0.059_amm
2.057
2.075
CytoBridge (Ours)
0.015
0.013-aax5
0.014:13:
0.021 3000
0.018
0.043-oux
0.038 30038
0,058
ccoel
8

---

Mouse Blood Hematopoiesis To demonstrate the scalability of CytoBridge to high- high-dimensional data, we adopt the mouse hematopoiesis dataset (Weinreb al. 2020) which includes 49,302 cells with lineage tracing data collected at three time points. We use PCA to reduce the dimensions to 50 and serve as the input CytoBridge. The dataset comprises diverse cell states and demonstrates pronounced cell division. Consequently, accurately modeling both cellular dynamics and growth rates is critical for reliable inference cell fate. As shown in Table 2, CytoBridge outperforms other Mate-of-the-an methods in distribution matching, highlighting CytoBridge's s capabilities of capturing transition patterns (Fig. 3(a).. Besides, evidenced by the TMV metric, CytoBridge able to recover the increase in cell numbers by learning the growth rate (Fig. 3(b)). The regions with high learned growth rates correspond to the hematopoietic stem cell populations. This is biologically consistent with the lineage tracing barcode results (Sha et al. 2024). The score function learned by CytoBridge indicates the existence of multiple attractors, which may lead to different cell fates (Fig. 3(c)). To further demonstrate the impact of cell-cell interactions the transition of cells, we computed the correlation of each cell's drift and interacting force. As shown in Fig. 3(d), the learned cell-cell interactions may promote early-stage cell differentiation and inhibit later-stage cell differentiation. Some additional results can be found in Appendix B.2.
(+)
Learnt Velocity
(b)
Learnt Growth Rate
high
(c)
Leartt Score
Correlation
Neutrophil
Homotyte
Engairoid
Lympiheedd
Meg
Figure 3: Application in mouse blood hematopoiesis data (a = 0.1), visualized in UMAP space. (a) The overall velocity learned by CytoBridge. (b) The growth rates learned by CytoBridge. (c) The score function learned by CytoBridge at f 2. (d) The correlation velocity and interacting forces.
Embryoid Body, Pancreatic B-cell differentiation and A549 EMT We further evaluated Cyto- Bridge on the Embryoid Body single-cell data which consists of 16,819 cells collected at five time points (Moon al. 2019). We used PCA reduce the dimensions to 50 and serve as the input of CytoBridge. As shown in Table 9, CytoBridge generally outperforms other methods in distribution matching and maintains mass-matching results comparable to other unbalanced algorithms. We plot- ted the learned velocity and score in Fig. 6, indicating the different cell fates learned by CytoBridge. The correlation of each cell's drift and interacting force shown Fig. 6(d) indicates that the learned cell-cell interactions may promote cell differentiation while inhibiting cell differentiation of some outliers. also consider the 3D-cultured in vitro pancreatic f-cell differentiation dataset (Veres et al. 2019) and A549 lung cancer cell line epithelial-messnchyym by TGFB1 (Cook and Vanderhyden 2020). Interestingly, we find that the interaction in the A549 cell line EMT process may be very weak. The detailed results can be found in Appendices B.3 to B.5.
transition (EMT) dataset induced

---

Table 2: Wasserstein distance (W)) and Total Mass Variation (TMV) of predictions at different time points across five runs on mouse hematopoiesis data (a 0.1). We show the mean value with one standard deviation, where bold indicates our algorithm as the best among all algorithms.
t = 1
t = 2
Model
Wi
TMV
Wi
TMV
SF2M (Tong al. 2024b)
8,217 10001
2.231 =0.00
11.086
=0.00
5.399=00.00
Meta FM (Atanackovic et al. 2025)
8.545 20,000
2.231 =0:00
10.313
5.399
MMFM (Rohbeck et al. 2025)
7.647:000
2.231 20.00
10.156=0000
5.399920.00
Metric FM (Kapusniak et al. 2024)
7.788 0.0000
2.23 I 0,000
11.449 10.000
5,399
UOT-FM (Eyring al. 2024)
2.114:00000
0.100 0.000
9.170
0.118 0.000
MIOFlow (Huguet et al. 2022)
6,313 3.313100 1:0000
2.23 0,0000
6.746
5,399 1..90000000
uAM (Neklyudov et al. 2023)
7.537 1:0000
2.875
9.762.10.00
5.670
UDSB (Pariset et al. 2023)
10.687 +ouoss
0.282
13,477 +0.05
3.010
TIGON (Sha et al. 2024)
6.140
=.234
6.973 0.000
2.083 0.000
DeepRUOT (Zhang et al. 2025a)
6.052 +0002
0.200
6.757 00000
0.260=o..M
CytoBridge (Ours)
6.013 20002
0.208
6.644=0.011
0.078=n.01
Extension to Spatiotemporal Transcriptomies To demonstrate CytoBridge's applicability in modeling cellular interactions with explicit physical proximity, we applied CytoBridge to zebrafish spatiotemporal transcriptomics dataset (Liu et al. 2022a), using slices from 5.25 hpf and 10 hpf as input. We evaluated the performance of CytoBridge on the task of reconstructing the dynamics of cell spatial migration, space and gene expression space respectively. As shown in Table 13, CytoBridge outperforms other methods in both tasks. The results are visualized in Fig. 9. Furthermore, downstream interpretability analysis of the learned interactions identified biologically relevant pathways crucial for zebrafish development. Detailed results be found in Appendix B.6. The preliminary application of CytoBridge to spatiotemporal transcriptomics demonstrates our framework's ' potential in modeling spatially resolved data.
well gene expression, with separate velocity and interactions for physical
Conclusion
We have introduced CytoBridge for learning unbalanced stochastic mean-field dynamics from time- series snapshot data. To tackle the interacting particle system inference from temporal snapshots, CytoBridge transforms the SDE constraints into ODE leveraging Fisher regularization for more effi- cient simulation in training. We have demonstrated the effectiveness of our method on both synthetic gene regulatory networks and single-cell RNA- data, showing its promising performance. Overall, CytoBridge provides a unified framework for generative modeling of time-series transcriptomics data, enabling more robust and realistic inference of underlying biological dynamics.
Limitations and Further Directions While CytoBridge offers valuable insights into incorporating cell-cell interaction with unbalanced stochastic dynamics, several aspects could benefit from further exploration. Firstly, CytoBridge minimizes the upper bound of the energy term in the UMFSB problem. Directly optimizing the original UMFSB formulation remains important question. Potential solutions may involve leveraging neural SDEs or the integration-by--are strategy proposed by (Zhang al. 2025a) as viable approaches. Secondly, the current neural network parameterization of cellular interaction terms is based on techniques used to reduce degrees of freedom, such as RBF expansion. Future work could explore sparse representation methods to replace it for improved expressive power. Thirdly, the training of CytoBridge involves multiple stages. Simplifying the training process by incorporating optimality conditions (e.g., HJB equations) presents a promising research direction (Sun et al. 2025). Furthermore, the current modeling of cellular interactions have not incorporated certain biological priors such as ligand-receptor information. We plan to explore this direction in future work to further refine our approach. Finally, extending the concept of flow matching to this context and developing simulation-fre training methods for stochastic, growth/death, and interaction dynamics could also advance the CytoBridge's applicability.
10

---

Acknowledgments and Disclosure of Funding
We thank Dr. Zhenfu Wang (PKU), Prof. Chao Tang (PKU) and Prof. Qing Nie (UCI) for their insightful discussions. This work was supported by the National Key R&D Program of China (No. 2021YFA 1003300 to T.L.), National Natural Science Foundation of China (NSFC No. 12288101 to T.L.& & P.Z., and Central Universities, Peking University. We acknowledge the support from the High-performance Computing Platform of Peking University for computation.
T2321001 to P.Z.) and The Fundamental Research Funds for the
References
[1] Michael S Albergo, Nicholas M Boffi, and Eric Hunden-Eijndes "Stochastic interpolants: unifying framework for flows and diffusions". In: arXiv preprint (oXiv:2303.08791 (2023).
[2] Axel A Almet, Yuan-Chen Tsai, Momoko Watanabe, and Qing Nie. "Inferring pattern-drivin intercellular flows from single-cell and spatial transcriptomics". In: Nature Methods 21.10 (2024), pp. 1806-1817.
[3] Lazar Atanackovic, Xi Zhang, Brandon Amos, Mathieu Blanchette, Leo Lee, Yoshua Bengio, Alexander Tong, and Kirill Neklyudov. "Meta Flow Matching: Integrating Vector Fields the Wasserstein Manifold". In: The Thirteenth International Conference Learning Representations. 2025.
[4] Nathalie Ayi and Nastassia Pouradier Duteil. "Mean-field and graph limits for collective dynamics models with time-varying weights". In: Journal of Differential Equations 299
(2021), pp. 65-110.
[5] Julio Backhoff, Giovanni Conforti, Ivan Gentil, and Christian Léonard. "The mean field Schrödinger problem: ergodic behavior, entropy estimates and functional inequalities". In: Probability Theory and Related Fields 178 (2020), pp. 475-530.
[6] Aymeric Baradat and Hugo Lavenant. "Regularized unbalanced optimal transport as entropy minimization with respect to branching brownian motion". In: arXiv preprint 1XXi:21/1100166 (2021).
[7] Simon Batzner, Albert Musaclian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E Smidt, and Boris Kozinsky. "E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials", In: Nature communi-- cations 13.1 (2022), 2453.
[8] Immanuel Pora,, José A Carrillo, and Sondre T Galtung. "Mean field limit for one dimensional opinion dynamics with Coulomb interaction and time dependent weights". In: Nonlinear Analysis 240 (2024), p. 113462.
[9] Jean-Davia Benamou and Yann Brenier. "A computational fluid mechanics solution to mass transfer problem". In: Numerische Mathematik 84.3 (2000),
the pp. 375-393.
[10] Volker Bergen, Marius Lange, Stefan Peidli, F Alexander Wolf, and Fabian Theis. "Gen- eralizing RNA velocity to transient cell states through dynamical modeling". In: Nature biotechnology 38.12 12 (2020), pp. 1408-1414.
[11] Charlotte Bunne, Ya-Ping Hsieh, Marco Cuturi, and Andreas Krause. "The schrödinger bridge between gaussian measures has closed form". " In: International Conference on Artificial Intelligence and Statisties. PMLR. 2023, pp. 5802-5833.
[12] Charlotte Bunne, Geoffrey Schiebinger, Andreas Krause, Aviy Regev, and Marco Cuturi. "Optimal transport for single-cell and spatial omics". In: Nature Reviews Methods Primers 4.1 (2024). p. 58.
[13] Zixuan Cang. Yanxiang Zhao, Axel A Almet, Adam Stabell, Raul Ramos, Maksim V Plikus, Scott X Atwood, and Qing Nie. "Screening cell-cell communication in spatial transcriptomics via collective optimal transport". In: Nature methods 20.2 (2023), pp. 218-228.
[14] Louis-Pierre Chaintron and Antoine Diez. "Propagation of chaos: review of models, methods and applications. II. Applications". In: arXiv preprint arXiv:2106. 14872 (2021).
[15] Ao Chen, Sha Liao, Mengnan Cheng, Kailong Ma, Liang Wu, Yiwei Lai, Xiaojie Qiu, Jin Yang. Jiangshan Xu, Shijie Hao, et al. "Spatiotemporal transcriptomic atlas of mouse organogenesis using DNA panttall-patterned arrays". In: Cell 185.10 (2022a), pp. 1777- 1792.
11

---

[16] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. "Neural ordinary differential equations". In: Advances in neural information processing systems 31 (2018).
[17] Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou. "Likelihood Training of Schrödinger Bridge using Forward-Backward Backward SDEs Theory". In: International Conference on Learning Representations. 2022b.
[18] Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. "The most likely evolution of diffusing and vanishing particles: Schrodinger bridges with unbalanced marginals". In: SIAM Journal on Control and Optimization 60.4 (2022c), pp. 2016-2039.
[19] Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard. "An inter- polating distance between optimal transport and Fisher-Rao metrics". In: Foundations of Computational Mathematics 18 (2018a), pp. 1-44.
[20] Lenaic Chizat, Gabriel Peyré. Bernhard Schmitzer, and François-Xavies Vialard. "Unbal- anced optimal transport: Dynamic and Kantorovich formulations". In: Journal of Functional Analysis 274.11 (2018b), pp. 3090-3123.
[21] Lénaic Chizat, Stephen Zhang, Matthieu Heitz, and Geoffrey Schiebinger. "Trajectory infer- ence via mean-field langevin in path space". In: Advances in Neural Information Processing Systems 35 (2022), pp. 16731-16742
[22] David P Cook and Barbara C Vanderhyden. "Context specificity of the EMT transcriptional response". In: Nature communications 11.1 (2020), p. 2142.
[23] Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. "Diffusion schrödinger bridge with applications to score-based generative modeling". In: Advances in Neural Information Processing Systems 34 (2021), pp. 17695-17709.
[24] Jun Ding, Nadav Sharon, and Ziv Bar-Joseph. 'Temporal modelling using single-cell tran- scriptomies". In: Nature Reviews Genetics 23.6 (2022), pp. 355-368.
[25]
Nastassia Pouradier Duteil. "MeaN-FIELD LIMIT OF COLLECTIVE DYNAMICS WITH WEIGHTS", In: Networks and Heterogeneous Media 17.2 (2022), pp. 129-
161.
[26] Luca Eyring, Dominik Klein, Théo Uscidda, Giovanni Palla, Niki Kilbertus, Zeynep Akata, and Fabian Theis. "Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation". R In: The Twelfih International Conference on Learning Representations. 2024,
[27] Xuanrui Feng and Zhenfu Wang. "Quantitative Propagation of Chaos for 2D Viscous Vortex Model with General Circulations on the Whole Space". In: arXiv preprint arXiv:2411. (2024).
[28] Jean Feydy, Thibault Séjourné, François-Xavies Vialard, Shun-ichi Amari, Alain Trouve and Gabriel Peyré. "Interpolating between Optimal Transport and MMD using Sinkhorn Divergences". In: The 22nd International Conference on Artificial Intelligence and Statistics. 2019, pp. 2681-2690.
[29] Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z Alaya, Aurélie Boisbumon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, et al. "Pot: Python optimal transport". In: Journal of Machine Learning Research 22.78 (2021),
pp. 1-8.
[30]
Nicolas Fournier, Maxime Hauray, and Stéphane Mischler. "Propagation of chaos for the 2D viscous vortex model". In: Journal of the European Mathematical Society 16.7 (2014),
pp. 1423-1466.
[31] Yifeng Fu, Hong Qu, Dacheng Qu, and Min Zhao. "Trajectory Inference with Cell-Cell Inter- actions (TICCI): intercellular communication improves the accuracy of trajectory inference methods". In: Bioinformatics 41.2 (2025), btaf027.
[32] Adam Gayoso, Philipp Weiler, Mohammad Lotfollahi. Dominik Klein, Justin Hong, Aaron Streets, Fabian Sareets,Fabian1 Theis, and Nir Yosef. "Deep generative modeling of transcriptional dynamics for RNA velocity analysis in single cells". In: Nature methods 21.1 (2024), pp. 50-59.
[33] Ivan Gentil, Christian Léonard, and Luigia Ripani. "About the analogy between optimal transport and minimal entropy". In: Annales de la Faculté des sciences de Toulouse: Mathé- matiques. Vol. 26. 3. 2017, pp. 569-600.
[34] Anming Gu, Edward Chien, and Kristjan Greenewald. "Partially Observed Trajectory In- ference using Optimal Transport and a Dynamics Prior". In: The Thirteenth International Conference on Learning Representations. 2025.
12

---

[35] Jiequn Han, Ruimeng Hu, and Jihao Long. "Learning nigh-dimensional -lasoo stochastic differential equations with general distribution dependence", In: SIAM Journal on Numerical Analysis 62.1 (2024), pp. 1-24.
[36] Matthieu Heitz, Yujia Ma, Sharvaj Kubal, and Geoffrey Schiebinger. "Spatial Transcriptomies Brings New Challenges and Opportunities for Trajectory Inference". In: Annual Review of Biomedical Data Science 8 (2024).
[37]
Camilo Hernández and Ludovic Tangpi. "Propagation of chaos for mean field Schrödinger problems". In: SIAM Journal on Control and Optimization 63.1 (2025), pp. 112-150.
[38] Jonathan Ho, Ajay Jain, and Pieter Abbeel. "Denoising diffusion probabilistic models". In: Advances in neural information processing systems 33 (2020), pp. 6840-6851.
[39] Yuanfei Huang, Chengyu Liu, and Xiang Zhou. "Lévy Score Function and Score-Based Partiele Algorithm for Nonlinear Lévy-Fokker-Planck Equations". In: arXiv preprint arXiv:24/2. (2024).
[40] Guillaume Huguet, Daniel Sumner Magruder, Alexander Tong, Oluwadamilola Fasina, Manik Kuchroo, Guy Wolf, and Smita Krishnaswamy. "Manifold interpolating aprimol-transport flows for trajectory inference". In: Advances in neural information processing systems 35 (2022), pp. 29705-29718.
[41]
Pierre-Fimannel Jabin and Zhenfu Wang. "Mean field limit for stochastic particle systems". In: Active Particles, Volume I: Advances in Theory, Models, and Applications (2017), pp. 379 402.
[42]
Qi Jiang and Lin Wan. "A physics-infrmmed neural SDE network for learning cellular dynamics from time-series scRNA-seq data". In: Bioinformatics 40 (2024), pp. -i127.
ISSN: 1367-4811.
[43] Shi Jin and Lei Li. "On the mean field limit of the random batch method for interacting particle systems". In: Science China Mathematics (2022), pp. 1-34.
[44] Shi Jin, Lei Li, and Jian-Guo Liu. "Random batch methods (RBM) for interacting particle systems", In: Journal of Computational Physics 400 (2020), p. 108877.
[45] Shi Jin, Lei Li, and Jian-Guo Liu. "Convergence of the random batch method for interacting particles with disparate species and weights". In: SIAM Journal Numerical Analysis 59.2 (2021), pp. 746-768.
[46] Kacper Kapusniak, Peter Potaptchik, Teodora Reu, Leo Zhang, Alexander Tong, Michael Bronstein, Joey Bose, and Francesco Di Giovanni. "Metric flow matching for smooth interpo- lations on the data manifold". In: Advances in Neural Information Processing Systems 37 (2024), pp. (35011-135044
[47] Diederik Kingma and Max Welling. "Auto-Encoding Variational Bayes". In: arXiv preprint 13Xiv: /4 (2013).
[48] Dominik Klein, Giovanni Palla, Marius Lange, Michal Klein, Zoe Piran, Manuel Gander, Laetitia MlengPPpaaxanthon Michael Sterr, Lama Saber, Changying Jing, et al. "Mapping cells through time and space with moscot". In: Nature (2025), pp. 1-11.
[49] Takeshi Koshizuka and Issei Sato. "Neural Lagrangian Schrödinger Bridge: Diffusion Mod- eling for Population Dynamics". In: The Eleventh International Conference on Learning Representations. 2023.
[50] Hugo Lavenant, Stephen Zhang, Young-Hoon Kim, Geoffrey Schiebinger, et al. "Toward a mathematical theory of trajectory inference". In: The Annals of Applied Probability 34.1A
(2024), pp. 428-500.
[51] Christian Léonard. "A survey of the Schrödinger problem and some of its connections with optimal transport". In: Discrete and Continuous Dynamical Systems-Series 34.4 (2014), pp. 1533-1574.
[52] Jingyuan Li and Wei Liu. "Solving McKean-Vlasor Equation by deep learning particle method". In: arXiv preprint arXiv:2501. .00780 (2025).
[53] Lingxiao Li, Samuel Hurault, and Justin M Solomon. "Self-consistent velocity matching of probability flows". In: Advances in Neural Information Processing Systems 36 (2023), pp. 57038-57057.
[54] Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hame Maximilian Nickel, and Matthew Le. "Flow Matching for Generative Modeling". In: The Eleventh International Conference on Learning Representations. 2023.
13

---

[55] Chang Liu, Rui Li, Young Li, Xiumei Lin, Kaichen Zhao, Qun Liu, Shuowen Wang, Xueqian Yang, Xuyang Shi. Yuting Ma. al. "Spatiotemporal mapping of gene expression landscapes and developmental trajectories during zebrafish embryogenesis". In: Developmental cell 57.10 (2022a), pp. 1284-1298.
[56] Guan-Horng Liu, Tianrong Chen, Oswin So, and Evangelos Theodorou. "Deep generalized schrödinger bridge". In: Advances in Neural Information Processing Systems 35 (2022b),
pp. 9374-9388.
[57] Guan-Horng Liu, Yaron Lipman, Maximilian Nickel, Brian Karrer, Evangelos Theodorou, and Ricky T. Q. Chen. "Generalized Schrödinger Bridge Matching". In: The Twelfth International Conference on Learning Representations. 2024.
[58] Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos A Theodorou, Weili Nie, and Anima Anandkumar. "B2SB: Image-to-Image Schrödinger Bridge". In: arXiv preprint 2XXi:23322.05822 (2023).
[59] Jianfeng Lu, Yue Wu, and Yang Xiang. "Score-based transport modeling for mean-field fokker-planck equations". In: Journal of Computational Physics 503 (2024), p. 112859.
[60] Leland McInnes, John Healy, and James Melville. "Umap: Uniform manifold approximation and projection for dimension reduction". In: arXiv preprint arXiv: 03426 (2018).
[61] Kevin R Moon, David Van Dijk, Zheng Wang, Scott Gigante, Daniel B Burkhardt, William Chen, Kristina Yim, Antonia van den Elzen, Matthew Hirn, Ronald R Coifman, et al. "Visualizing structure and transitions in high-dimensional biological data". In: Nature biotechnology 37.12 (2019), pp. 1482-1492.
[62] Kirill Neklyudov, Rob Brekelmans, Daniel Severo, and Alireza Makhzani. "Action matching: Learning stochastic dynamics from samples". In: International conference machine learning. PMI.R. 2023, pp. 25858-25889,
[63] Kirill Neklyudov, Rob Brekelmans, Alexander Tong, Lazar Atanackovic, Qiang Liu, and Alireza Makhzani. "A Computational Framework for Solving Wasserstein Lagrangian Flows".
In: Forty-first International Conference on Machine Learning. 2024.
[64] Alessandro Palma, Till Richter, Hanyi Zhang, Manuel Lubetzki, Alexander Tong, Andrea Dittadi, and Fabian Theis. "Multi-Modal and Multi-Attribute Generation of Single Cells with CFGen".. In: The Thirteenth International Conference on Learning Representations. 2025.
[65] Mattco Pariset, Ya-Ping Hsich, Charlotte Bunne, Andreas Krause, and Valentin De Bor- toli. "Unbalanced Diffusion Schrödinger Bridge". In: ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems. 2023.
[66] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary Vitito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. "Automatic differentiation in pytorch". In: (2017).
[67] Qiangwei Peng. Peijie Zhou, and Tiejun Li. "stVCR: Reconstructing spatio-temporal dynam- ics of cell development using optimal transport". In: bioRxiv (2024), pp. 2024-06.
[68] Katarina Petrović, Lazar Atanackovic, Kacper Kapusniak, Michael M. Bronstein, Joey Bose, and Alexander Tong. "Curly Flow Matching for Learning Non-gradient Field Dynamics". In: Learning Meaningful Representations of Life (LMRL) Workshop at ICLR 2025. 2025.
[69] Chen Qiao and Yuanhua Huang. "Representation learning of RNA velocity reveals ro- bust cell transitions".. In: Proceedings of the National Acudemy of Sciences 118.49 (2021), e2 e21058591 18.
[70] George Rapakoulias, Ali Reza Pedram, and Panagiotis Tsiotras. "Steering Large Agent Populations using Mean-Field Schrodinger Bridges with Gaussian Mixture Models". In: arXiv preprint arXiv:2503.23303 (2025).
[71] Martin Rohbeck, Charlotte Bunne, Edward De Brouwer, Jan-Christian Huetter, Anne Biton, Kelvin Y. Chen, Aviv Regev, and Romain Lopez. "Modeling Complex System Dynamics with Flow Matching Across Time and Conditions". In: The Thirteenth International Conference on Learning Representations. 2025.
[72] Lars Ruthotto, Stanley Osher, Wuchen Li, Levon Nurbekyan, and Samy Wu Fung. "A machine learning framework for solving high-ding figh-dimensional mean field game and mean field control problems". In: Proceedings of the National Academy of Sciences 117.17 (2020), pp. 9183-9193.
14

---

[73] Geoffrey Schiebinger. Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, al. "Optimal- transport analy- sis of single-cell gene expression identifies developmental trajectories in reprogramning", In: Cell 176.4 (2019), pp. 928-943.
[74] Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexan- dre Tkatchenko, and -obur-t Müller. "Schnet: A continuous-filtt convolutional neural network for modeling quantum interactions". In: Advances in neural information processing systems 30 (2017).
[75] Yutong Sha, Yuchi Qiu, Peijie Zhou, and Qing Nie. "Reconstructing growth and dynamic trajectories from single-cell transcriptomics data". In: Nature Machine Intelligence 6.1 (2024), pp. 25-39.
[76] Zebang Shen and Zhenfu Wang. centropy-iisipattinn informed neural network for mekean- vlasov type pdes". In: Advances in Neural Information Processing Systems 36 (2023), pp. 59227-59238.
[77] Zebang Shen. Zhenfu Wang, Satyen Kale, Alejandro Ribeiro, Amin Karbasi, and Hamed Hassani. confscoonss of the fokker planck equation". In: Conference on Learning Theory. PMLR. 2022. pp. 817-841.
[78] Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. "Diffusion Schrödinger bridge matching". In: Advances in Neural Information Processing Systems 36 (2024).
[79]
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. "Deep unsu- pervised learning using nonequilibrium thermodynamics" In: International conference on machine learning. PMLR. 2015, pp. 2256-2265.
[80] Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. "Score- Generative Modeling through Stochastic Differential Equations". In: International Conference on Learning Representations. 2021.
[81] Kelly Street, Davide Risso, Russell B Fletcher, Diya Das, John Ngai, Nir Yosef, Elizabeth Purdom, and Sandrine Dudoit. "Slingshot: cell lineage and pseudotime inference for single- cell transcriptomics". In: BMC genomics 19.1 (2018), p. 477.
[82] Yuhao Sun, Zhenyi Zhang, Zihan Wang, Tiejun Li, and Peijie Zhou. "Variational Regularized Unbalanced Optimal Transport: Single Network, Least Action". In: Advances in Neural Information Processing Systems (2025).
[83] Alejandro
Paul Bertin, Stefan Bauer, Hananch Aliee, Yoshua Bengio, and Fabian J Theis. "Causal machine learning for single-cell genomics". In: Nature Genetics (2025), pp. 1-12.
[84] Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Pector-Brooks, Guy Wolf, and Yoshua Bengio. "Improving and generalizing flow-based generative models with minibatch optimal transport". In: Transactions on Machine Learning Research (2024a). Expert Certification. ISSN: 2835-8856.
[85] Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. "Tra- jectorynet: A dynamic optimal transport network for modeling cellular dynamics". In: Inter- national conference machine learning. PMLR. 2020. pp. 2526-9536.
[86] Alexander Tong, Manik Kuchroo, Shabarni Gupta, Aarthi Venkat, Beatriz San Juan, Laura Rangel, Brandon Zhu, John G Lock, Christine L Chaffer, and Smita Krishnaswamy. "Learning transcriptional and regulatory dynamics driving cancer cell plasticity using neural ODE-based optimal transport". In: bioRxiv (2023), pp. 2023-03 03.
[87] Alexander Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio. "Simulation- Free Schrödinger Bridges via Score and Flow Matching". In: International Conference on Artificial Intelligence and Statistics. PMLR. 2024b, pp. 1279-1287.
[88] Cole Trapnell, Davide Cacchiarelli, and Xiaojie Qiu. "Monocle: Cell counting, differential expression, and trajectory analysis for single-cell RNA-Seq experiments". In: Bioconductor (2017).
[89] Blias Ventre, Aden Forrow, Nitya Gadhiwala, Parijat Chakraborty, Omer Angel, and Geoffrey Schiebinger. "Trajectory inference for branching SDE model of cell differentiation" In: arXiv preprini urXiv:2307.0766 (2023).
15

---

[90] Adrian Veres, Aubrey L Faust, Henry L Bushnell, Elise N Engquist. Jennifer Hyoje-Ryu Kenty, George Harb, Yeh-Chuil Poh, Elad Sintov, Mads Gürtler, Felicia W Pagliuca, et al. "Charting cellular identity during human in vitro B-cell differentiation". In: Nature 569.7756 (2019), pp. 368-373.
[91] Wei Wan, Yuejin Zhang. Chenglong Bao, Bin Dong, and Zuoqiang Shi. "A scalable deep learning approach for solving high-dimensional dynamic optimal transport". In: Journal on Scientific Computing 45.4 (2023), B544-B563. B563.
[92] Dongyi Wang, Yuanwei Jiang, Zhenyi Zhang, Xiang Gu, Peijie Zhou, and Jian Sun. "Joint Velocity-Growth Flow Matching for Single-Cell Dynamics Modeling". In: Advances in Neural Information Processing Systems (2025).
[93] Yusong Wang, Tong Wang, Shaoning Li, Xinheng He, Mingyu Li, Zun Wang, Nanning Zheng, Bin Shao, and TTe-Yann Liu. "'Ehhancing geometrie representations for molecules with equivariant vector-scalar interactive message passing". In: Nature Communications 15.1 (2024), p. 313.
[94] Caleb Weinreb, Alejo Modriguuzz-Fraticllle Fernando D Camargo, and Allon M Klein. "Lin- eage tracing on transcriptional landscapes links state fate during differentiation". In: Science 367.6479 (2020), eaaw3381.
[95] F Alexander Wolf, Fiona K Hamey, Mireya Plass, Jordi Solana, Joakim S Dahlin, Berthold Göttgens, Nikolaus Rajewsky, Lukas Simon, and Fabian Theis. "PAGA: graph abstraction reconciles clustering with trajectory inference through topology preserving map of single cells". In: Genome biology 20.1 (2019). p. 59.
[96] Haoming Yang. Ali Hasan, Yuting Ng, and Vahid Tarokh. "Neural -lasen-Vlasov processes: distributional dependence in diffusion processes". In: International Conference Artificial Intelligence and Statistics. PMLR. 2024, pp. 262-270.
[97] Liu Yang, Constantinos Daskalakis, and George E. Karniadakis. "Generative ensemble re- gression: Learning particle dynamics from observations of ensembles with physics-informed deep generative models". In: SIAM Journal on Scientific Computing 44.1 (2022). B80-B99. [98] Maosheng Yang. "Topological Schrödinger Bridge Matching". In: The Thirteenth Interna- tional Conference on Learning Representations. 2025.
[99] Grace Hui Ting Yco, Sachit D Saksena, and David K Gifford. "Generative modeling of single- cell time series with PRESCIENT enables prediction of cell trajectories with interventions". In: Nature communications 12.1 (2021), p. 3222.
[100] Yuning You, Ruida Zhou, and Yang Shen. "Correlational Lagrangian Schrödinger Bridge: Learning Dynamics with Population-Level Regularization". In: arXiv preprini arXiv:2402. 10227 (2024).
[101] Jiaqi Zhang, Erica Larschan, Jeremy Bigness, and Ritambhara Singh. "scNODE: generative model for temporal single cell transcriptomic data prediction". In: Bioinformatics S0ppup
ment_ 2 (2024a), pp. ii 1114-iil554 54. ISSN: 1367-4811.
[102] Peng Zhang. Ting Gao, Jin Guo, and Jinqiao Duan. "Action Functional as Early Warning Indicator in the Space of Probability Mcasures". In: arXiv preprint arXiv:2403. 10400 (2024b).
[103] Zhenyi Zhang, Tiejun Li, and Peijie Zhou. "Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport". In: The Thirteenth International Confer- ence on Learning Representations. 2025a.
[104] Zhenyi Zhang, Yuhao Sun, Qiangwei Peng, Tiejun Li, and Peijie Zhou. "Integrating Dynam- ical Systems Modeling with Spatiotemporal scRNA-Seg Data Analysis". In: Entropy 27.5
(2025b). ISSN: 1099-4300.
[105] Zhenyi Zhang, Zihan Wang, Yuhao Sun, Jiantao Shen, Qiangwei Peng, Tiejun Li, and Peijic Zhou. "Deciphering cell-fate trajectories using spatiotemporal single-cell transcriptomic data". In: Preprint (2025c).
[106] Linqi Zhou, Aaron Lou, Samar Khanna, and Stefano Ermon. "Denoising Diffusion Bridge Models". In: The Twelfth international Conference on Learning Representations, 2024a.
[107] Mo Zhou, Stanley Osher, and Wuchen Li. "Score-based Neural Ordinary Differential Equa- tions for Computing Mean Field Control Problems". In: arXiv preprint arXiv:2409.1644 (2024b).
[108] Qunxi Zhu, Bolin Zhao, Jingdong Zhang, Peiyang Li, and Wei Lin. "Governing equation discovery of a complex system from snapshots". In: arXiv preprint arXiv:2410. 16694 (2024).
16

---

A Training Details
A.1 Training CytoBridge
The training of CytoBridge involved training vo. 90. 86 and hg. These networks were initialized after the pre-traing phase, and the overall training process involved minimizing the following loss:
C Ennegg + CRosons + CF
The calculation of each loss component involved the numerical solution of temporal integrals and ODEs, which was performed using the Neural ODE solver (Chen et al. 2018). The gradients of the loss function with respect to the parameters for Vø. 90. So and t were derived through neural ODE computations, and these networks were optimized using Pytorch (Paszke et al. 2017). To calculate the reconstruction loss during training, we utilized the implementation of Sinkhorn algorithm provided by (Feydy et al. 2019).
A.2 Training Initial Log Density Function through Score Matching
Conditional Flow Matching (CFM) is used to leam initial log density. First, sample pairs (Xo, X1) are chosen from an optimal transport plan (Xx., x1), and Brownian bridges are constructed between them. We initially assume (t) to be constant.
The log density is matched with these bridges, where p(x, t (xo, x1)) = N N(;; (x1 + (1 x)xo, o - t)) and its gradient is Vx log p(x, &. (xo.x1)) =
txi +11+(1-1)x- -X for E [0, 1]. The
x=x(1-1)
neural network sæ (x, t) is used to approximate o p(x,(x,0) t) with weighting function Ag. The unsupervised loss Lus is:
Laus = X 200 -vVx logp(x,t)]
The corresponding CFM loss, facore: is:
1 V log,(x,1 / (xo, x1))
Loor = Rqq A V (x) 2
where Q' = (1. u(0, xxx @ p(xx1 1| (Xo; x1) By taking the weighting function as:
(t (1
The CFM loss can be converted to:
Seorr = (t) xx t) + e 11.
where E ~ (O, I). This formulation is computationally more tractable.
A.3 Modeling Interaction Potential
Inspired by the design of machine learning force fields in physics (Schütt et al. 2017; Batzner et al. 2022; Wang et al. 2024), we model the cell interaction network by expanding distances d, between cells and j using radial basis functions (RBFs), An exponential transformation is first applied to the raw distance drj. set of K RBF features, ex (diy), are then computed based on this exponentially scaled distance:
(d(4)) = cxp(-34 (cxp(-dr)) (cp((3)(e - 2)
The RBF centers μ4 are initialized by uniformly discretizing the exponentially transformed distance interval that corresponds to original distances from 0 to a predefined cutoff deutott. The width parameters 3% are initialized as the inverse square of term proportional to the average spread of these centers in the transformed exponential scale. The RBF expansion allows our model to learn interactions between different cells by encoding interaction distances across multiple scales, rather than enforcing interactions only between similar cells.
These expanded features, forming vector e(d,;) = [ei (dz)).... ex ex(dx))] (d,))l,, are subsequently fed into a multi-layer neural network (NN) with its own set of trainable parameters, NNN, to predict the interaction potential क. This relationship can be expressed as:
(xx,,,) xj) = NN(e(d,)::NNN
17

---

The interaction weight function (xx;,x) is defined as:
{o
if dig < deutoff
x(x
if dig 2 deunoff
A.4
Reconstruction Loss Function
The local mass matching loss Local Mass is designed to ensure that the distribution of weights among the sampled particles each time step aligns with the local density of the real data. Suppose we sample particles at different time points with batch size N, which is denoted as A, At given time step Êx, the local mass matching error Mg is computed for the N sampled particles. This error is defined as:
MR Σ
card (x:
72 Tio
Here, w:(ls) is the weight the i-th sampled particle at time (s, ny denotes the number of cells in the original dataset at time lg. The term card (hr (xx())) represents the number of sampled data points from As that are mapped the i-th sampled particle (f(x) via the mapping /gg This mapping hy assigns each real data point A₁ its closest sampled particle Ås. Essentially, the formula measures the squared difference between the particle's 's current weight and the proportion of real data points it represents. Thus, provides a fine-graines guidance on the weights of particles, Moreover, under the condition that M, = 0, it follows that Σi (xx = * Eie card (he (xx))) nu n noo Therefore the local matching loss also encourages the alignment of the total mass. The local mass matching loss Liucal Mass then calculated by summing these errors over time steps from == to T 1: Local Mass Ek=1 M₁-
Besides, global mass matching loss is further adopted to align the total number of cells at different time points during the training phase. Specifically, the global mass error term Gg is defined as:
Gk = Σ w. (tx)))
Tig
n2
Here, The global mass matching loss Litoball Mass is then calculated by summing these errors over all time steps from k = 1 to T 1: CGtobal Mass = TT GR.
The optimal transport loss is computed Lor (I/N,1/N; 1/N,..., denotes the uniform distribution of sampled points Aş at time ty with batch size N, w(tx) = (wi (t&), w2((x),... wN (tx))/ EN wi denotes the predicted weight distri- bution of particles at time tg.
Σ (y ((+)). Here w
A.5 Energy Loss Function
To simplify the computation, we adopt upper bound of the energy for training purposes:
/ 2 (vx(((t))
Lenery -Exorppo
| /V
+ at (ge) wa
B Additional Results
B.1 Synthetic Gene Regulatory Network
Dynamics By combining these interactions with the three-gene regulatory network, the system dynamics is defined as follows:
18

---

dxi
a1(x))² +3
N
1 + (x1) + az (x²)2 + + - x + 71 N
Σ
W Pi.j.
dt
j
dxp
ax +
82X2 + 1122t
kij V2
di
1+ + az (x2))) + 73 (Xg)² +
N j
dx3
a3(X3))
1+ (xa3) 6z 78
Σ jfij=1
K V
di
N
Here, (t) denotes the gene expressions of the i cell at time t, L. represents ( -X), while and 3 represent the strengths of self-activation, inhibition, and external stimulus respectively. The interaction weight function is defined based pre-defined threshold deutoff. where krjj is to if di, is within the threshold, otherwise ka,j O. The parameters 6; describe the rates of gene degradation, and T7&E represents stochastic influences via additive white noise. The probability of
x3
cell division is associated with X2 expression, given by the formula g 0gg T+X2 When cell
divides, new cells are generated with independent random perturbations nu NO, 1) for each gene around the gene expression profile (X1(t), (t)),, Xx())) of the parent cell. Hyper-parameter detailed in Table 3. The initial cell population is drawn independently from two normal distributions, ([2,, 0.2, 0j, 0.1) and ([OO 0, 2). 0.1). At each time step, any negative expression values are to 0.
Table 3: Simulation parameters on gene regulatory network.
Parameter
Value
Description
0.5
Strength of self-activation for X Strength of inhibition by X3 on X Strength of self-activation for X2
711
0.5
0g2
I
722
I
Strength of inhibition by X3 on X₂
03
1
Strength of self-activation for X3
73
10
Half- half-saturation constant for inhibition terms
61
0.4
Degradation rate for X1
82
0.4
Degradation rate for X₂ Degradation rate for X3 Noise intensity for X₂ Noise intensity for X2 Noise intensity for Xa
83
0.4
1/1
0.05
172
0.05
TT8
0.01
Md
0.014
Noise intensity for cell perturbations External signal activating X1 and X₂ Threshold of interaction weight function
3
1
deutoff
0.5
de
0.1
Equilibrium distance of the Lennard-Jona potential
F Finar
100
The upper bound of the magnitude of forces
dt
1
Time step size
Time Points
[0, 8, 14, 24, 32]
Time points at which data is recorded
Hold-one-out Evaluations To evaluate CytoBridge's capability of recovering distributions at unseen time points, we conducted hold-one-out experiments on synthetic gene regulatory data with attractive interactions. Specifically, intermediate time points were individually left out during training and subsequently recovered during evaluation. Note that the default setting of UDSB only supports data with three time points inputs, we only compare with the performance of other methods. As shown Table 4, CytoBridge performed best on all held-out time points, indicating its ability to infer reasonable trajectories from snapshots.
Lennard-Jones Interaction Potential The Cennard-Jones-lig interaction potential lets cells tend to prevent others from being too similar while attracting cells with distinct gene expressions. The Ø
19

---

Table 4: Wasserstein distance (Wi) of predictions at held-out time points across five runs on synthetic gene regulatory data with attractive interactions (a 0.05). We show the mean value with one standard deviation, where bold indicates the best among all algorithms.
t
t= 2
= 3
Model
Wi
Wi
Wi
SF2M (Tong et al. 2024b)
0.184
0.241 =00100
0.507
Meta FM (Atanackovic et al. 2025)
0.272=0000
0.293=n.00
0.344
MMFM (Rohbeck al. 2025)
0.476=0.:00
0.466=20000
1.215
Metric FM (Kapusniak et al. 2024)
0,188 0,000
0.498 00000
0.630
UOT-FM (Eyring et al. 2024)
0.204
0.155 0.000
0.136.10000
MIOFlow (Huguet et al. 2022)
0,225 0,000
0,270
0.234-10.00
uAM (Neklyudov al. 2023)
0,600
0.975
0.000
1.243 10000000
TIGON (Sha et al. 2024)
0.254
0.2 4 0.000
0.178+a.00
DeepRUOT (Zhang et al. 2025a)
0.184
0.086
0.004
0.079 20003
CytoBridge (Ours)
0.182:t0.00
0.064
0.001
0.043 40002
defined as:
||x (1x->))) (1=0(
de
12
de
6
( - y) =4
yll
||x
yl
where de represents the distance at which the repulsive and attractive forces balance each other. Moreover, to avoid the singularity of the potential function, we clipped the magnitude of forces to pre-defined value Fmoz- As shown in Fig. 4, CytoBridge is able to reproduce the Lennard-Jone interaction potential with only information from snapshots provided, leading to correct transition and change of variance. We also conducted quantitative evaluations compared with DeepRUOT, in which the cellular interactions are neglected. As shown in Table 5, CytoBridge consistently outperforms DeepRUOT across all time points, underscoring the importance of explicitly modeling cellular interactions.
Table 5: Wasserstein distance (Wi) of predictions for DeepRUOT and CytoBridge (Ours) on synthetic gene regulatory data with Lennard- -oones Interaction Potential (a = 0.05).
1
t
f = 3
4
Model
W2
Wi
Wi
Wi
DeepRUOT
0.042=m000 0.036400000
0.049
0.054
CytoBridge (Ours)
0.030+t0.0
0.030 =.030=0.01
0.027
0.031 +000s
(a)
(b)
Intercation Potentaa 40)
00) Nm
Figure 4: Results of synthetic gene regulatory data with Lennard-Joner Interaction Potential. (a) The dynamics learned by CytoBridge. (b) The learned interaction potential.
No Interaction Despite the good performance of CytoBridge on the synthetic gene model with different types of interactions, it yet remains unknown whether our method can help identify the
20

---

cases where the ground truth dynamics itself does not involve cellular interactions. Therefore, we further conduct experiments on the synthetic gene model without interactions involved by setting (xx y) 0. As shown in Table 6, explicitly incorporating cellular interactions does not result in the improvement of performance in distribution matching at most time points. The main reason for this phenomenon is that it may be hard for neural network to directly infer all-zero outputs. The numerical error resulting from our interaction network may lead to perturbations at certain time points and thereby exhibit poorer performance than methods that neglect the interaction term. We further analyze the learned interaction forces acting on each cell. As shown in Fig. 5, compared with forces learned from the dynamics with explicit interactions, forces learned from dynamics with interaction exhibit notably smaller magnitudes. We compute the correlation of learned velocity and forces of each cell to examine whether the learned forces from data without interaction may exhibit certain patterns related to the transitions of cells. Moran's I, which is a statistic that measures spatial autocorrelation, indicating the degree which nearby locations have similar values, is utilized to identify whether these patterns exist. Moran's 's calculated from data with ground truth interactions is 0.514, while that calculated from data without interactions is 0.281, which is notably smaller. This may indicate that instead of explainable correlated patterns between transition and interaction, interacting forces learned from data without interaction exhibit random patterns to some extent. Biologically, it remains challenging to precisely know whether and how cells interact with each other. Based on this difficulty, we hope that our preceding observations can help explain the performance of our method when applied to real-world biological datasets. Furthermore, we aim for this analysis to provide valuable biological insight into the presence and nature of cell-cell interactions.
Table 6: Wasserstein distance (Wi) of predictions for DeepRUOT and CytoBridge (Ours) synthetic gene regulatory data without interaction potential (a = 0.05).
t =
t = 3
=
4
Model
W1
W1
Wi
Wi
DeepRUOT
0.035
0.054=0000
0.046:1002 0.044.2000
0.042.1000
CytoBridge (Ours)
0. 0.03610.001
0.053
0.05210003
(a)
4.0
20
155
10
Cene
Figure 5: Comparasions of interacting forces learned from (a) dynamics with Lennard-Jones interac- tion potential, (b) dynamics without interaction potential.
B.2 Mouse Blood Hematopoiesis
We obtained the dataset from (Weinreb et 2020). We used the cells profiled from in vitro cultured mouse hematopoietic progenitors, grown in conditions supporting multi-lineage differentiation. While not transplanted, the system retains extensive intercellular signaling, making it suitable for studying interaction-dricn transition dynamics. The original gene expression space was projected to 50-dimensional subspace using PCA. We summarized the results of hold-one-out evaluations in Table 7. The improvements observed in both Table 2 and Table 7 indicate that incorporating interaction terms indeed helps to recover the trajectories of cells. Moreover, the learned interacting forces exhibit significant correlation with the learned velocity, quantified by Moran's 's I of
21

---

0.757. Based on our preceding analysis, this observed pattern and the improvement in performance strongly suggest the presence of genuine interaction forces within this realworld data. To enhance interpretability, we identified top 200 genes influenced most significantly by cell-cell interactions. Subsequent enrichment analysis of these genes revealed key pathways associated with the biological processes, as shown in Table 8. Applying this to the mouse hematopoiesis dataset, we identified pathways related to positive regulation of leukocyte activation and cell-cell adhesion which are closely linked to hematopoiesis, interactions, and differentiation. We also incorporated Trajectory Inference with Cell-Cell Interactions (TICCI) (Fu et al. 2025) as a reference method. By applying TICCI to the mouse hematopoietic dataset, it identified ligand-recptor pair Lgals9-Cd455 which is closely related to the regulation of ' cell activation, further provides biological interpretation for enriched pathways. Moreover, calculated the gradients of growth network with respect to genes to identify key genes that contribute to growth dynamics, including Meis Nfkb1, and Rap Meis1 regulates hematopoietic stem cell proliferation and self-renewal, Nfkbl promotes cell cycle entry and survival as a signaling hub, and Raplb drives cell division via pathways like MAPK, consistent with established biological knowledge. These findings demonstrate CytoBridge's ability to uncover biologically relevant mechanisms.
Table 7: Wasserstein distance (Wi) of predictions at held-out time point across five runs on mouse hematopoiesis data (o 0.1). We show the mean value with one standard deviation, where bold indicates the best among all algorithms.
Model
W1
SF2M (Tong et al. 2024b)
8. 3.646_n.00
Meta FM (Atanackovic al. 2025)
10.82 1
MMFM (Rohbeck et al. 2025)
8.263
Metric FM (Kapusniak et al. 2024)
7.753
UOT-FM (Eyring et al. 2024)
9.332.10.000
MIOFlow (Huguet et al. 2022)
7.779 0.000
uAM (Neklyudov et al. 2023)
9.157=00.0 0.000
TIGON (Sha et al. 2024)
8.402 0.000
DeepRUOT (Zhang et al. 2025a)
6,868 2.003
CytoBridge (Ours)
6.847
Table 8: Enriched pathways from interactions of mouse hematopoiesis data, showing the adjusted p-value and gene count for each term.
Pathway
p.adjust
Count
regulation of T cell activation positive regulation of cell activation
3.38 x 10 12
23
1.42 x 10-11 11
23
positive regulation of cell-cell adhesion
5.81 x 10-11
20
positive regulation of leukocyte activation
1.90 10-10
21
leukocyte cell-cell adhesion lymphocyte differentiation
3.18 x 10-10 10
21
1.04 x 10-9
21
B.3 Embryoid Body
We use the same dataset in (Moon al. 2019; Tong al. 2020), which consists of 16,819 cells collected at five time points. The cells are from human embryoid bodies (EBs), formed in 3D culture by spontaneous aggregation of stem cells over a 27-day differentiation time course. The experimental setup mimics early developmental conditions, where diverse cell types emerge and are expected to interact through signaling and spatial organization. We projected the original gene expression space to 50 dimensions using PCA. Hold-out evaluations were performed at four distinct time points, and our method consistently yielded the best results in these tests (Table 10). However, the observed performance gain, while positive, was less pronounced compared to the results on the mouse hematopoiesis data. Correspondingly, the learned interacting forces for this dataset exhibit correlation pattern in the cell-state landscape with a Moran's 's I of 0.450. This moderate positive
22

---

Moran's I suggests the presence of cell-state cunsition-relatee interaction effects within this data, though potentially weaker or less organized than those inferred for the lineage-constraing systems such as mouse hematopoiesis data.
Table 9:
Wasserstein distance (Wi) and Total Mass Variation (TMV) of predictions at different time points across five runs on embryoid body data (a 0.1). We show the mean value with one standard deviation, where bold indicates the best among all algorithms.
(
t 2
3
4
Model
TMV
TMV
TMV
34
TMV
SF2M (Tong et al. 2024b)
9.146
0,748
10.882-.vx
0.377
11.650
0.539
12.154.a007
0.399
Meta FM (Annackovic et al. 2025)
2.497
0.748g
11.054_0mmm
0.377
=1.567=mm
0.539
11. 11.487
0.399 me
MMFM (Robbookk
2025)
2.124
0.748
10:474..ww
0.377
11.022:uum
0.539
11,480aaee
a.
Metric FM (Kapusniak al. 2024)
8.506 Louae
0.748 ceae
9.795 aureo
0.377
10.621
0.539
12.042 aaoeo
0.3991cax
UOT (Eyring et al. 2024)
9. 20000
0.054
10.982-mmm
0.078
11.584
0.014
12. 12:
0.092 ac
MIOFlaw (Huguet al. 2022)
8.447
0.748.cax
9,229
0.377
9.436
0.539
10.123
BAM (Neklyudov et 2025)
12.315
14.996-21m
1.323
15.0852a00
1.536
10.407 =nen
1.396 se
UDSB (Buinc et 2023)
11 11.9%3:une
0.429
14,009 .com
0.005
14.65 14.656_aums
0.166
15. 15.884220022
0.029 ac
TIGON (Sha et al. 2024)
8.433 uoudo
0.118
9.275
0.030...3
9,802
0,276
10.148 anoio
0,141
DeepRUOT (Zhang al. 2025a)
2.15
0,050
9.034gorex
0.161
9,369
0.005.4eam
9.772
CytoBridge (Ours)
8.159 200
0.002.tamm
9.027 gnues
0.057
9,351 Looee
0.175
9.750mmmm
0.022
Table 10: Wasserstein distance (Wi) of predictions at held-out time points across five runs on embryoid body data (o = 0.1). We show the mean value with one standard deviation, where bold indicates the best among all algorithms.
1=1 1
t= 2
1 1= 3
Model
W1
W1
W1
SF2M (Tong et al. 2024b)
+0000
11,276 10002
11.380
Meta FM (Atanackovic et al. 2025)
10.504
11.478 10000
11.660
MMFM (Rohbeck et al. 2025)
10.23910.00 10000
11.46910000
11.930.:0.000
Metric FM (Kapusniak et al. 2024)
9.672+a000
11.041 +0000
11.466
UOT-FM (Eyring et al. 2024)
10.366-+0.000
13.583
15.858 +0.000
MIOFlow (Huguet et al. 2022)
10.684+40.00
11.755
10.440##.no
uAM (Neklyudov et al. 2023)
12.857 1.0000
15.743
2..000
TIGON (Sha et al. 2024)
11.19920.00
11.207-20000
10.833=0.00m
DeepRUOT (Zhang et al. 2025a)
9.628
10.382.2000
10.215
CytoBridge (Ours)
9.62622001
10.333
10.201 =20.007
B.4 Pancreatic 3-cell Differentiation Data
We evaluated our method on the dataset from (Veres et al. 2019), which contains 51,274 cells collected across eight time points from human pluripotent stem cells differentiating toward pancreatic 3-like cells in 3D suspension culture. The original gene expression space was projected to 30 dimensions using PCA. As shown in Fig. 7, CytoBridge is able to infer the velocity and growth rate, while identifying attractors. When comparing our approach, which explicitly models cellular interactions, to DeepRUOT, an algorithm that does not, we observed improved performance in 5 out of the 7 tested time points. Furthermore, the learned interacting forces for this dataset exhibited moderate correlation distribution pattern with inferred state-transition velocity, quantified by Moran's 's I of 0.590. Overall, the performance gains across most tested time points, coupled with the observed moderate patterns in the inferred forces, may suggest the presence of interaction forces within this dataset. The correlation between forces and velocities indicates that cells may tend to prevent others from differentiation at early stages while promoting differentiation at later time points.
B.5
EMT Data
We use the dataset from (Sha et al. 2024; Cook and Vanderhyden 2020), derived from A549 cancer cells undergoing TGFB1-ine -induced maithelial-mesenhhyme four distinct time points. The cells were cultured in standard 2D monolayers and treated with TGFB1 over several days. Although EMT is coordinated process in vivo involving paracrine and
transition (EMT). This dataset comprises
23

---

[2]
Learnt Velocity
(b)
Learnt Growth Rate
high
low
(c)
Learnt Score
(d)
Correlation
Figure 6: Application in embryoid body data (o = 0.1), visualized in PHATE space. (a) The overall velocity learned by CytoBridge. (b) The growth rates learned by CytoBridge. (c) The score function learned by CytoBridge at f. = 4. (d) The correlation of velocity and interacting forces.
Learnt Velocity
[b)
Learnt Growth Rate
high
Leamt Score
(d]
Correlation
Figure 7: Application in pancreatic 3-cell differentiation data (a = 0.1), visualized in UMAP space. (a) The overall velocity learned by CytoBridge. (b) The growth rates learned by CytoBridge. (c) The score function learned by CytoBridge 7. (d) The correlation velocity and interacting forces.
24

---

Table 11: Wasserstein distance (W;) of predictions for DeepRUOT and CytoBridge (Ours) at different time points on pancreatic 8-cell differentiation data (a 0.1).
t
3
4
5
1 6
Model
W2
W/
Wj
W/
W4
Wi
W4
DecpRUOT
8.0447
7.6301 =wam
8.0064
7.9018 g
8.3977 sqmme
7.8346 7.8343=0a00
CytoBridge (Ours)
8.0771+o
7.6299-00002
100066
7.9018=0mu7
8.3974=201
contact-dependent signaling, this in vitro system mimics EMT largely cell-autonomous response to an external stimulus. Despite the fact that CytoBridge still is able to infer the velocities, the growth of cells, and different cell fates (Fig. 8). When comparing our method, which explicitly models interaction terms, to approaches like DeepRUOT that do not, we observed no improvement in performance on this dataset. Furthermore, the distribution pattern of the learned interacting forces for this dataset was very disorganized, yielding Moran's of 0.040. This indicates that the inferred forces are largely random compared to the transition velocity direction, suggesting lack of significant or organized intercellular interactions for state-transition The absence of both performance gain and significant correlation in inferred forces suggests that transitions in this setting primarily driven by direct transcriptional responses rather than intercellular signaling, which is consistent with the biological experiment setup.
Table 12: Wasserstein distance (W;) of predictions for DeepRUOT and CytoBridge (Ours) at different time points EMT data (a 0.05).
= 1
2
t
3
Model
Wi
W1
Wi
DeepRUOT
0.239
0.253
0.261
CytoBridge (Ours)
0.240
0.259
0.269
[8]
Velscity
Learnt Growth Rate
low
(4)
Correlation
Figure 8: Application in EMT data (a = 0.05), visualized in PCA space. (a) The overall velocity learned by CytoBridge. (b) The growth rates leamed by CytoBridge. (c) The score function learned by CytoBridge at / 4. (d) The correlation of velocity and interacting forces.
25

---

B.6 Zebrafish Spatiotemporal Data
We adopt the Zebrafish Embryogenesis Spatiotemporal Transcriptomic Atlas (ZESTA) (Liu al. 2022a), created using the Stereo-seq spatial transcriptomics technology (Chen et al. 2022a). The dataset provides a high-resolution map of gene expression in zebrafish embryos across six critical time points within the first 24 hours of development, Among the six time points, we selected 5.25 hpf and 10 hpf as input. To align spatial coordinates hetween these two time points, we adopt the rigid body transformation invariant optimal transport. Then, we projected the original gene expression space to 50 dimensions using PCA. To model cellular interactions in both physical space and gene expression space, we transformed both the distances in physical and gene expression space into separate RBF features, and concatenated them as inputs to the interaction potential. Thus, the effects of cellular interactions are achieved by calculating the gradients of interaction potential with respect to distances in physical and gene expression space respectively. The reconstruction loss is calculated in both physical space and gene expression space. We compared the W1 distances in the physical space and gene expression space between CytoBridge and other methods. As shown in Table 13, CytoBridge achieves better performance over other methods in both physical space and gene expression space. We also visualize the predicted cell states in physical space and gene expression space in Fig. 9. To further interpret the biological effects of learned cellular interactions the development process of zebrafish, we identified top 200 genes influenced most significantly by interactions. We identified pathways related to somite development which are critical to zebrafish embryonic development based on the enrichment analysis, as shown in Table 14. These findings align with known biological processes, showing the framework's potential in spatially resolved data.
Table 13:
Wasserstein distance (Wi) of predictions on zebrafish embryogenesis data. We report metrics in physical space (denoted as "Space") ) and gene expression space (denoted as 'Gene'). Bold indicates the best result.
Model
Space
Gene
SF2M (Tong et al. 2024b)
0.265
5.423
Meta FM (Atanackovic al. 2025)
0.268
5.413
MMFM (Rohbeck et al. 2025)
0.247
5.208
Metric FM (Kapusniak et al. 2024)
0.273
5.366
UOT-FM (Eyring et al. 2024)
0.227
5.173
MIOflow (Huguet al. 2022)
0.263
4.720
uAM (Neklyudov et al. 2023)
0.177
6.499
TIGON (Sha et al. 2024)
0.352
4.979
DeepRUOT (Zhang et al. 2025a)
0.261
4.745
CytoBridge (Ours)
0.035
4.712
Table 14: Enriched pathways from zebrafish embryogenesis data, showing the adjusted p-value and gene count for each term.
Pathway
p.adjust
Count
nucleolus
5.76 x 10
12
somite development
1.82 x 10-4
10
lipid transport
1.82 x 10-4
12
gastrulation
4.15 x 10
12
somitogenesis
4.15 x 10
8
B.7 Ablation Studies
We conducted ablation studies on the synthetic gene regulatory data with attractive interactions to demonstrate the effectiveness of CytoBridge's 's different components. We first note that without the interaction term, our method reduces to the framework of DeepRUOT. Compared with RUT, CytoBridge performs better in all these metrics, underscoring the effectiveness of explicitly consider- ing cell-cell interactions. We then examine the impact of growth term on our algorithm. First, we set the growth term g to zero, and observe that omitting the growth term will lead to poorer performance
26

---

(a)
T=0
(b)
T=1
(c)
(d)
Booker
Pooterm Your
Segmeraal
Schct u01
Figure 9: Application in zebrafish embryogenesis data (a 0.05). (a) The spatial coordinates of cells sampled 0. (b) The predicted spatial coordinates of cells 1 by CytoBridge (c) The predicted gene expression by CytoBridge in the PCA space. (d) The overall velocity learned by CytoBridge gene expression space.
in distribution matching, which may result from the false transition of balanced Schrödinger Bridge solvers. We then evaluate the impact of Lass By setting the weight of mass loss to zero, observe that although taking the growth term into account indeed helps eliminate the false transition compared to the results without growth, it still falls short in capturing the changes in total mass, evidenced by the TMV metric. Therefore, it is important to incorporate the LMsss term to match the mass changes. Furthermore, examine the impact of Fokker-Planck Constraint. By setting the weight of LFP to zero, we also observed an overall drop in performance. As the Fokker-Planee constraint is necessary restrain the relationship of our different networks, omitting it will prevent the score function from correctly reflecting the distributions of cells.
Next, we investigate the role of pretraining. Without the pretraining phase, CytoBridge exhibits significant drop in performance across all time points, with considerably higher distribution matching metrics and less accurate mass matching at later time points. This indicates that pretraining is crucial for initializing the networks effectively and achieving stable overall training. Omitting the score matching phase used initialize the matching. This suggests that the score matching phase is vital for effectively training the score network. Finally, analyzing the model without the main end-to-end training (relying solely pretraining, denoted as "w/o training"), "), we observe performance drop in distribution matching across all time points compared to the full CytoBridge model, and large discrepancies in TMV later stages. This confirms that while the pretraining phase provides a beneficial initialization, the complete end-to-end training procedure is indispensable for achieving CytoBridge's superior results.
function also leads to poorer performance in distribution
Table 15: Wasserstein distance (W1) and Total Mass Variation (TMV) of predictions with different settings different time points across five on synthetic gene regulatory data with attractive interactions (a 0.05). We show the mean value with one standard deviation.
3
Model
Wi
TMV
Wi
TMV
W;
TMV
W;
TMV
CytoBridge interaction (DespRUOT)
0.059
0.075
CytoBridge wio growth
a.
0.080
0.175
0.250
0.343
0.515
0.480
0.930sr
CytoBridge wio Crus
0.016axx
0.9422mm
0.037arm
0.165aras
0.0464s39
0.352asm
01047
0.689
CytoRridge
6.014:
0.01325m
0.016
0.030
0.022
0.048
0.06
CytoBridge wio pretraining
0-426.a30
0.016..am
0.915
0.035
1,1951x30
0.0441xa4
1.451 1oaM
1.86
CytoBridge wño score matching
0.025 saxx
0.011
0.025
0.024.caaa
0.034.x300
0.049
0.039 au0
0.084 :cas)
CytoBridgo trining
0.33
0.012.aas
0.02911a0
0.02
0.067
1030
0.050.ea25
0.108
0.174.om
CytoBridge (Ours)
0.015
0.014::ax
0.021
27

---

C Experimental Details
C.1 Evaluation Metries
We evaluate the 1-Wasserstein distance (W)), which used to evaluate the similarity between inferred distributions and true distributions, and Total Mass Variation (TMV), which is used to evaluate whether the inferred dynamics is able to reflect the growth of cells, on the synthetic gene data and real-world single-cell data. The metries are defined as follows:
/ ||= - 2112d1
(i 4) = min (EII(P-))
where p and represent empirical distributions. The Wi is calculated using the Python Optimal Transport library (POT) (Flamary et al. 2021).
TMV( = Σ (tx)
T2k
where (tx) represents the weight particle at time lus and vı denotes the number of cells in the original dataset at time fx-
For all datasets, models were trained using all available time points and were evaluated using W1 and TMV. For synthetic gene regulatory networks with attractive interactions, mouse hematopoiesis, and embryoid body dataset, additional experiments with one-time point held out were conducted. Wi used to evaluate the performance at held-out time points. We evaluate our method by applying the learned dynamics to all initial data points to generate trajectory and their weights for subsequent time points starting from initial weights (;) 1/mo. Next, compute the weighted Wi and TMV between the generated data and the real data based on the inferred weights. We ran the simulation five times to calculate the mean value and the standard deviation. To evaluate the performance of DeepRUOT (Zhang et al. 2025a), we deploy the same procedure as ours but without interaction. To evaluate the performance of TIGON (Sha et al. 2024), reimplemented their method to avoid certain instabilities. As for uAM (Neklyudov et al. 2023), we deploy its default parameter settings. TIGON and uAM are evaluated by simulating the dynamics of weighted particles.
To ensure fair comparison with other methods, we used their default settings for datasets featured in their original papers. For all other datasets, we adjusted the models' network sizes to ensure comparable parameter counts and tuned their training epochs and learning rates accordingly. To evaluate the performance of SF2M (Tong et al. 2024b), we keep the diffusion coefficient the same as ours while maintaining other parameters as defaults for fair comparison. The weights of the inferred particles are set to uniform distribution to evaluate as the growth term is not considered. To evaluate the performance of UDSB (Pariset et al. 2023), as its default setting only involves three-time points, we use samples at 0, 2, 4 from synthetic gene data and embryoid body data inputs.
C.2 Hyperparameters Selection and Loss Weighting
The experiments were performed on a shared gigh-performance computing cluster with NVIDIA A100 GPU and 128 CPU cores. As we aim to make our algorithm universally applicable to different types of biological data, most of the hyperparameters are kept identical across different datasets. while those varying are mainly related to the scale of datasets. Specifically, for the modeling of interaction potential, we set the number of RBF kernels to 8 across different datasets. For real- world scRNA-seq data, the threshold demott is set no lower than the largest distance between cells in specific datasets so that all pairs of cells are involved in interacting with each other. For the zebrafish data, we set the threshold deutoff in the physical space, where cells interact with neighboring cells in physical space. We conducted experiments on the effect of different values of deutoff in Table 16.
As shown, we found that the performance was quite robust with respect to different choices of deutolf unless the cutoff is set too small, which may potentially lead to inadequate modeling of cellular interactions. To simulate ODEs with random batch methods, we also examined the choices of number of particles p within one group in Table 17. Generally, increasing p slightly improves performance on evaluation metrics by providing a more accurate approximation of the interaction term. However, gains diminish at higher p values, while computational costs increase. Thus, we selected p = 16 as the default, balancing robust performance with computational efficiency.
28

---

Table 16: Wasserstein distance (Wi) and Total Mass Variation (TMV) of predictions with different
deutoft values different time points. The table shows the mean value with one standard deviation.
deunott
t
1
2
t
3
4
W1
TMV
W1
TMV
W1
TMV
W4
TMV
0.1
a,080
0.010=a00s
0.1
0,032+a016
0.141 0.141=00m +0.00
0.050=0042
0.167 #2000
0.077 #0043
0.3
0,015+am:
0.013+amu
0.016+no0
0.024+a.22
0.026+a.0
+.038+0044
0.042#2004
0.052+o06
0,5
0.015+am0
0.013+aa00
0.014+no0:
0,021 +a.24
0.018+a.mm
0.043 +0044
0,038 +2000
0.058 +oc6l
0.7
0.013+amm:
0.013+mm0:
0.024+aa02
+.024
0.047 +00001
0.037
+00000
0.069 #0/066
1.0
0.014::0m:
0.012:a000
0.014
0,025 a20023
0.026
0.039 #:004
0.043 0:004
0.057 #0066
2.0
0.018
0.025 a.01
0.023
0.042
0.046=0001
0.049
0.072
Table 17: Wasserstein distance (W1) across different time points for varying values of parameter p.
P
=
1
== 2
== 3
t
2
0.025
0.033
0.048
0.094
0.017
0.023
0.025
0.043
8
0.016
0.019
0.022
0.035
16
0.015
0.014
0.018
0.038
32
0.015
0.016
0.020
0.030
64
0.014
0.014
0.024
0.030
Moreover, we also compared the results with and without the random batch method in Table 18. Results show that the full model (without RBM) and the RBM model yield comparable performance, with the full model slightly better at certain time points. However, RBM significantly enhances computational efficiency, reducing the interaction term's 's complexity from O(N2) OppN). The tables demonstrate substantial reductions in memory and inference time with RBM. Given its comparable accuracy and scalability for large-scale single-cell datasets, RBM is justified and essential component of our framework.
Table 18: Performance comparison simulation data with and without RBM.
t
1
t t= 2
t t= 3
t t= 4
Time (s)
Memory (GB)
Method
Wi
Wi
Wi
Wi
w/o RBM
0.015
0.013
0.023
0.029
0.568
22.1
p = 16
0.015
0.014
0.018
0.038
0.115
1.7
For our training procedure, the parameters differ only in the number of training epochs. As real-world ScRNA-seq data mainly involves large numbers of cells, it typically will require more iterations for our model to converge. Increasing the number of training epochs has few adverse effect, allowing users to use a default of (500, 100, 500) epochs for larger datasets. For other ryper-parameters we will then provide guideline on their choices. In the pre-training phase, we first need to provide suitable initialization for the velocity network and the growth network, which involves selecting the parameters Am and Aa in LRccoos Am LMass + XaCoT- We set Aa to 1 and Xm to 0.01 in order to encourage the transition to match the observed distribution while maintaining the unbalanced effect. We empirically found that lowering the mass loss weighting during pre-training improves distribution matching performance. Here we only adopt the local mass matching loss without restricting the exact number of cells in order to learn general growth patterns. Subsequently, we set Xm to zero to initialize the interaction network. By doing so, the interaction network can be stably trained in order to refine the variance of distributions. The parameter selections of pre-training procedure is summarized in Table 19, where the arrow notation (->) represents the adjustment of hyperparameters during two stages of our pre-trate mreitraining phase.
During the training phase, as the four networks have been reasonably initialized, these networks can be trained together stably. We set Am and X to 1, while adding the global mass matching term to encourage our growth network to match the exact number of cells at different time points. d in L LEosrg + CRocons Cpp is set to better match the distributions. Xy and Au in Lyp
29

---

set to align the score network to match the density. The specific choices of these parameters are summarized in Table 19. We utilize the set of parameters consistently across different datasets.
Table 19: Parameter Settings for Different Datasets Across Two Training Stages (Synthetic gene, mouse hematopoiesis, embryoid body, pancreatic 3-cell differentiation, EMT, Zebrafish).
Parameter
Synthetic gene
mouse bernllopiccss
embryoid body
pancicalic 9-cell difierentiatio
EMT
Zebrliish
Pre-Training Phase
(Am Epochs)
(001, 200) (040. 100)
10.01. 5,500)
00.01. 500) (0.0, 100)
(0.01. 500) (0.0, 100)
(0.11 100)
(0.01. 5,50)
(0.0, 100)
(0.0, 100)
(0.0. 100)
Training Phase
Am
x 10a
x
1
1x 10
x 10
x 10
x 104
10
1 x 10
1x 10
x 10
10*
10
dease Epochs
0.5
300
100
100
2
0.3
200
1000
500
500
200
500
C.3
Scalability and Computational Efficiency
We conducted an evaluation of the scalability and computational efficiency of CytoBridge on the embryoid body data by extending the input from 50 to 100 PCs. As shown in Table 20, CytoBridge achieves the lowest Wi distance across all time points, demonstrating that CytoBridge remains effective at higher dimensionality. Regarding the computational efficiency, CytoBridge maintains comparable to other neural ODE-based methods with a training time 11 minutes and peak GPU memory usage of 6.3 GB.
Table 20: Performance on 100D embryoid body data. We show the Wasserstein distance (W,) at four time points, alongside total runtime and peak memory usage. Bold indicates the best result.
1=1
1
=
==
3
Time
Memory (GB)
Model
W
Wi
Wi
Wi
SF2M (Tong al. 2024h)
11.333
12,982
13.718
14.945
4min 43s
0.7
Meta FM (Atanackovic et al. 2025)
11.699
13.398
14.037
14.727
Smin 50s
MMFM (Rohbeck et al. 2025)
13.150
14.135
14.441
14.907
4min 38s
0.6
Metric FM (Kapusniak et al. 2024)
10.806
12.348
13.622
16.801
2min 16s
4.2
UOT-FM (Fyring et al. 2024)
10.757
12.799
13.761
15.657
3min 27s
0.6
uAM (Neklyudov al. 2023)
13.628
18.315
20.309
22.973
38s
1.5
MIOflow (Huguet et al. 2022)
11.387
12.331
11.905
12.908
6min 19s
TIGON (Sha et al. 2024)
10.547
12.926
13.897
14.535
7min 33s
0.7
DeepRUOT (Zhang ct al. 2025a)
10.226
11.110
11.544
12.424
10min 17s
3.7
CytoBridge (Ours)
10.217
11.070
11.505
12.368
11min 26s
6.3
C.4 Visualization
For the Mouse hematopoiesis and pancreatic 3-cell differentiation data, project them di- mension using UMAP (MeInnes et al. 2018) for visualization. For embryoid body data, we project them to 2 dimension using PHATE (Moon al. 2019). For EMT data, we use the first two PCs for visualization. The learned velocity, score and interaction are visualized using sc Velo (Bergen et al. 2020). Here, the velocity stands for the combination of the drift b(x, t) of corresponding SDE and interacting forces, which drives the process of the transition of cells. Specifically, se Velo projects high-dimensional vectors into a nower-dimensinn reflecting cell-to-cell transition probabilities; these probabilities are based on the cosine similarity between the target high-dimensional vector and the displacement vector to neighbors. Using this matrix and the cells' positions in chosen embedding, seVelo then estimates the embedded vector for each cell as the expected displacement, effectively summarizing the likely future state of the cell in the low-dht www-dimensional representation.
space by first computing transition matrix
30

---

D Mathematical Details
D.1 Proof of Theorem 4.1
Proof. From Definition 4.1 we obtain
Op -
[(((x.)) / k(x, -x)(x) p(x,1) + r²(t) Ap(x, + g0,
aL
- 2 [ pxx,) 1 ((xx y)ply, 1) dy p(x.,) ) + gp.
Using the change of variable v(x, b(x, ((x,1))-/2(()V)_. log lox (), we that it is equivalent y)l(y, dy) td p(x,t) 9p.
0p(x, t) Vx ((vxx:)) - J/r 1(xx (
Ot
Correspondingly, the integrand in the objective functional becomes
inf 1' J 1 + a'(t) (p.b.g,0)
8 V log p1l + 1 +5 log + all tav( p(x, l)dxdl.
(5)
D.2 Derivation of Proposition 5.1
We assume the following conditions hold.
Assumption D.1. The initial positions X3 are independently and identically distributed (i.i.d.) with a common density Pa(x) L' (RR) n L (R.). and the initial weights are set as w,(0) = 1 for all i = =,,..., N. The functions b, g. 6. and & are Lipschitz continuous and bounded: specifically, b and g are Lipschitz in x uniformly in I, k((x y) is Lipschitz in both arguments and bounded, and 2 5. is Lipschitz continuous and bounded. The diffusion coefficient o is continuous and bounded Ein w;({)xx; converges to a deterministic limit p(x,, L).
on [0, T. The empirical measure μda
The system in Proposition 5.1 satisfies argument N + co, i.e. N-T (t)
Jax f(x, t), t), where is an arbitrary test test function.
Remark D.I. The argument we assume here is related to the notions of "chaos" and " "propugation of chaos" " and il has a rich theory mathmatics (Jabin and Wang 2017; Chaintron and Diez 2021). To rigorously prove such an argument, refer some related works (Fournier al. 2014; Feng and Wang 2024; Duteil 2022; cen-Porat et al. 2024: Ayi and Duteil 2021).
We derive the macroscopic continuity equation from the microscopic particle dynamics using the weak formulation, employing test functions and taking the mean-fiel limit.
Define the weighted empirical measure as
Σ (1()8xx;
This measure encapsulates the distribution of particles along with their weights. Our goal is to show that, as N + oo, fu converges weakly to p(x, 1) dx, where satisfies the stated PDE.
Consider smooth test function yo Fd + IR in C (Rd) (i.e., infinitely differentiable with compact support). We examine the time evolution of the pairing
1 [ (0xx) (dx) = N
Σ (t))())))
To handle the stochasticity, we take the expectation and compute
N
d
w.(0)((3))
E
di
31

---

and then pass to the limit as N + oo to recover the weak form of the PDE.
For each particle i, the quantity wy (X)) evolves due both the deterministic weight dynamics and the stochastic position dynamics. Since wx (t) follows an ODE and x; follows an SDE, we apply Itô's product rule:
d (t))()) (xx) (t) + (t) (x())) dw;(t) dip(x;).
The weight evolution term is given by
dwi(l) - g X; (t) dt,
so
4p (xx;) (t, = g( (t) dt.
For the position evolution, apply Itô's formula to p(X;):
1
8
dxp(X)) (x)(X)) dx + 2
(Xi) (a())) dt,
duk
where the diffusion term arises from the Wiener process (with variance o² (t)). Substituting the SDE
dx;-
b(X: 1) N
kiy OOV
xx) de dl)(()dd;
J
we obtain
dx((x)) = V
[
b
1
Σ
o²(t) (x( dt. ki ",, x() de (L)
N
2
Thus.
Σ (.,, x/) di o²(1) A (X))
(:(()))((X)) ui)()
(x((X)) b
V dw;
Since (t) g(X;, (t); dt is deterministic (with no stochastic component), the cross variation (t) dxp(x)) = O. Combining these, the total differential
d (t)(())( = ,(xi, t) (t) (()))
b-
(t x() de
N - j4
+ (1) (2(() Ax - (())))))) .dW:
2
Summing over all N particles and dividing by N. we have
d dt N Σ wx(())((3))
N
g g(x) (t))))
N
(xt())
b
Σ
2t) xX)
J#1
o²(t) wtt
1 N
dW;
Ax(X))
(,tt)o(t)V
dt
Taking the expectation, the stochastic term vanishes since E V
(X.) dw: = 0. yielding
d
Σ m.(())(()))
(xx
di
b Σ jf
(1V x/)
+ (t(() (2( 2
AI((X))) )
32

---

As N - oo, we invoke the assumption. The empirical measure vua converges to p(x, dx. The interaction sum
N k 2,, x
approximates the mean-field interaction
J VV 2x; t) - dy.
By the regularity of di and V. 4, this approximation holds in expectation as N - cc. Thus, in the limit,
d di [ ((x))(xx))dxx
S go0 + p(x, ()Vxpx))
0 (2)
p(x, ()xp()) dx.
Rewrite the right-hand side using integration by parts, noting that yp has compact support (so boundary terms vanish). The drift term becomes
f pVp -Vxx-bxx- bdx - bax==
f 4v (bp)
The interaction term is
JoL pV.. (( s)p(y,)) dy dx
Jel xV [p
k(x,y)V s)p(y, dy dx.
Rd
The diffusion term is
Jau pxp = / pApdx,
and the growth term is
f dx.
Thus.
[ po dx 2d
- V- (bp) + V. (0/ k(x, N.x../ +
o²(() 2
=
40
Ap + gp dx.
di
Since this equality holds for all E C (RR",, the fundamental lemma of calculus of variations implies
-
Öp
bo-bbpppp k(x,y) ( dy o²(t) Ap + gpp.
Øt
which is the desired continuity equation.
E More Background on Trajectory Inference
In this section, provide more background on the trajectory inference task. In single-cell transcrip- tomics, methods vary depending on the data type:
Single Time-Point Snapshot Data: In this context, trajectory inference methods are broadly catego- rized into two groups The first type is pseudotime-based methods (Trapnell et al. 2017; Street et al. 2018; Wolf et al. 2019), which infer trajectories by ordering cells along pseudotime axis but often require specifying starting point, limiting their flexibility, The second type RNA velocity- methods (Bergen et al. 2020; Qiao and Huang 2021; Gayoso et al. 2024), which leverage spliced and
33

---

unspliced counts to estimate dynamics but are constrained by the need for such data. Both approaches infer dynamics from existing data without generating new cell states.
Multi Time-Point Snapshot Data: Our work addresses a distinct scenario involving single-cell sequencing data across multiple time points. Due the destructive nature of technology, this be framed as a generative modeling task, where dynamics are inferred from distributions at different time points and could be interpolated unseen time points. Once learned, these dynamics enable the generation of intermediate cell states. Methods like optimal transport, which have gained significant attention in computational systems biology and machine learning, are well-suited for this task (Zhang et al. 2025c).
Regarding evaluation and comparabilityy, generative approach, unlike pseudotime or RNA velocity methods, evaluates performance using metrics for generated distributions, which traditional methods cannot produce. This fundamental difference in data and modeling objectives, along with our hold-out experiments, renders direct comparisons with traditional trajectory inference methods infeasible. Zhang et al. 2025b discusses modeling approaches for different data types. We summarizes these distinctions in Table 21.
Table 21: Comparison of Trajectory Inference Methodologies
Method
Data
Key
Generative
Evaluation
Pseudotime
Single snapshot or merged snap- shots without us- ing time labels
Infers trajecto- ries by ordering
No
Pseudotime ac- curacy with true time labels
cells; requires
starting point
RNA Velocity
Single snap- shots with spliced/unsplcee
Estimates dy- No namics using RNA splicing kinetics
Velocity consis- tency
counts
Optimal Transport
Multi-time- point snapshots
Infers dynamics and generates
Yes
Generated distri- bution metrics
new cell states
F
Broader Impacts
Our work presents a new step forward in data-driven modeling of complex, dynamic biological systems by introducing the UMFSB framework and an associated deep learning methodology Cyto- Bridge, capable of explicitly accounting for cell-cell interactions alongside stochastic and unbalanced population effects. The enhanced ability to dissect intercellular communication within evolving cellular landscapes offers the potential for deeper mechanistic insights into fundamental biological processes such as organismal development, disease pathogenesis, and tissue regeneration. Moreover, approach could enable more accurate predictions of individual therapeutic responses, aid in the design of optimized combination therapies or cell-based treatments. By providing more accurate representation of these systems, our approach may accelerate the discovery of novel therapeutic targets, helpful for more effective healthcare interventions.
Although the prospective benefits for scientific understanding and biomedical application consider- able, the use of such predictive models also requires careful consideration of potential social impacts and risks. As with any data-driven approach, biases present in training datasets could be amplified by the model, potentially leading to biased outcomes applicd without rigorous check. Consequently, thorough validation across diverse conditions and systematic assessment of potential biases needed.
34

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction clearly state the claim. The theoretical portions and the algorithm of the claim are supported in Section 4 and Section 5; the claims on its performance are supported by Section of the main text.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations of the work in section Section
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-specfiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
35

---

Answer:
[Yes]
Justification: Proofs and assumptions are included in the main text and Appendix D.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Reproducibility is accomplished by providing the full description of the algorithm in the main text in addition to providing available code in a public repository /ttps://gitee http//gttttbtt//////// www/coo//aaaeoeee wwam///heniiiii zhang CytoBridge -ertoriidge-NuuIPS..
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
36

---

The authors should provide scripts to reproduce all experimental results for the new reproducible, they Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Code and data is available in a public repository https://github. com/ hatprriiiaannnttrrrriiiea IP..
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (uttps://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
/ httppsscccppbccccccceecce hepollic/giieeeee Codeeubbm ssi aasonnooicyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We specify the process of training and testing for all methods presented in the paper. See Appendix C.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification: The paper includes standard deviation (std) of the mean for all the results.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
37

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma enror bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide the description of the compute resources utilized with available computing clusters detailed in Appendix C.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: / :/tps:://nurrrr httpr/eurip...ee cc cctpppblic/Ethiccccuueeleellllll
Answer: [Yes]
Justification: We have all reviewed and confirmed this research conforms with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Broader impacts are discussed Appendix
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
38

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All the datasets used for analysis in this study are publicly available from the original source. They are properly credited, respected, mentioned and used. See Appendix B.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
39

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new datasets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
40

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The method development does not involve LLMs as any important, original, non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy (https:/ htt.::/p..../oooeeeeeeee///22550 for what should or should not be described.
41


Reviews:
[{"review_id": "e0kaXCbpAT", "reviewer": "Reviewer_idfK", "summary": "The authors tackle the challenging problem of learning the dynamics of interacting particles (cells) from sparsely time-resolved observations. Motivated by a body of literature from this domain, they formulate this problem as an Unbalanced Mean-field Schrödinger Bridge (UMFSB) problem, collectively enabling the modelling of stochastic dynamics, birth-death of cells, and cell-cell interactions. To approximate the UMFSB, the authors propose a novel deep learning-based approach they label as CytoBridge. Through experiments on both synthetic and real datasets, the authors demonstrate that CytoBridge outperforms its counterpart SOTA baseline methods for the trajectory inference task.", "questions": "- In line 19 of Algorithm 2, do you need to evaluate/simulate the integral during training? \n- I don't fully understand how to interpret the learned interaction potential results. More generally, how do the empirical results/plots validate that CytoBridge is learning the correct cell interactions? Could the authors please elaborate on this?\n- How does CytoBridge scale to higher dimensions? I believe 50 PCA components for the single-cell datasets is the largest dimensionality considered in this work? It would be useful to provide results for 100 PCs to observe this.\n- Regarding the Random Batch Method (RBM), ODE simulation is used to reduce the computational complexity of the ODE simulation with cellular interaction terms. What $p$ (which I understand is the size of the random batch) did you try? I see in Appendix C.2 $p = 16$. Did you try other values for $p$? How does performance change for different $p$'s? \n- Similarly, have you conducted an ablation to empirically test how the RBM inference method compares to standard inference? I understand theoretically they should lead to equivalent results, but it may be useful to confirm this empirically. \n\nMinor comments:\n\n- In Definition 3.2, $\\phi(\\cdot)$ (lowercase) is introduced but not used anywhere in the definition.\n- Figure 1 is never referenced in the text.", "limitations": "Yes, but in the appendix. I encourage the authors to use the extra page allotted for the final version of the manuscript to include the discussion of limitations (and future work) in the main text.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths:**\n\n- The paper is well written and clear. The authors do a good job of motivating the problem they are addressing and demonstrating a comprehensive understanding of the field / related work. \n- The problem addressed in this work, learning dynamics of interacting particles (cells), is clearly formulated and introduced as a mean field unbalanced schrödinger bridge problem. The authors do a good job going from the high-level problem -> mathematical definition of the problem -> how to solve it with machine learning -> empirical experiments. \n- The results are generally convincing, barring some drawbacks in comparisons and presentation (see weaknesses and questions). The proposed method, CytoBridge, outperforms all baselines for the task of multi-marginal trajectory inference on both synthetic and real datasets.\n\n**Weaknesses:**\n\n- The authors considered a wide range of trajectory inference baselines for comparison with CytoBridge, however, for the majority of the baselines, the default hyperparameters are used. I recognize that tuning all baselines is time-consuming and not realistic, but to conduct a truly fair comparison, I believe it is important to apply some additional work into searching for reasonable hyperparameter settings for the baselines. I think this is specifically important for methods that don't explicitly use (i.e. train their models on) the datasets considered in this work.\n    - In this vein, at the minimum, a simple thing to do is to make sure the networks that parameterize the vector fields are of comparable size in terms of the number of parameters, across all models/baselines. \n- The proposed method, CytoBridge, which approximates the formulated mean field unbalanced schrödinger bridge problem, requires the optimization of 4 neural networks -- which likely comes with its own challenges. Although the authors briefly mention this as a limitation in Appendix E, there seems to be limited discussion on the procedure for selecting/finding good hyperparameters (some is included in Appendix C.2). With many hyperparameters to select, this does not seem like a simple process. Some added discussion would be valuable. Especially given that a PINN loss is being used, which, from my understanding, can be difficult to optimize.\n    - I also would like to see a discussion of the challenges/effort of optimizing 4 neural networks for this method and how it relates to hyperparameter selection to some of the baselines.", "was_revised": true, "final_justification": "The authors did a good job of addressing my questions and concerns. Looking at their rebuttal and discussion with other reviewers as well, I see that the authors have added a _significant_ amount of additional experiments. As a consequence, I think the empirical quality of this paper will be improved. To add, the authors are trying to tackle a very difficult problem, learning the dynamics of cells and their accompanying interactions (or more generally, learning the dynamics of interacting particles). These problems are independently challenging, and I have assessed the quality of their contribution in light of this. Lastly, this paper provides a reasonable methodological/algorithmic component, which I believe is well-suited for NeurIPS. Hence, I have raised my score to an accept.", "rebuttal": "> Summary and Strengths\n- Thank you. We sincerely appreciate your recognition of the strengths in our paper. In this paper, we address the challenge of **jointly** modeling cell-cell interactions and unbalanced stochastic biological effects using multiple single-cell snapshots data. This makes our work an **early and essential** step in the integration of **interpretable cellular interactions** with **complex cell dynamics** within the optimal transport framework, offering both **great theoretical and practical potential** for exploring the **spatiotemporal dynamics and interactions of single cells**. \n\n> W1: Hyperparameter of baselines\n- Thank you. We fully agree that a rigorous and fair evaluation is essential. To address this, **we have made every effort to enhance the baseline comparisons by thoroughly tuning key baselines**. Specifically, we have implemented the following improvements：\n    - We have increased the network sizes of the baseline models to ensure their parameter counts are comparable to CytoBridge.\n    - We have extended the training epochs for these models beyond default settings to better suit different datasets.\n    - We have incorporated a learning rate scheduler into the training process to optimize performance.\n- These adjustments have **enhanced the performance of most baselines**. Nevertheless, our algorithm achieves improved performance over tuned baselines, indicating the robust contribution of CytoBridge. Here we present our tuned results ($\\mathcal{W}_1$) on the **mouse hematopoiesis** data as below, and **we will update all relevant tables in the manuscript with the results from these tuned baselines.**\n\n|Model (Mouse Data)|t=1|t=2\n|:---|:---|:---\n|SF2M|8.217|11.086\n|Meta FM|8.545|10.313\n|MMFM|7.647|10.156\n|MIOFlow|6.313|6.746\n|uAM|7.537|9.762\n|UDSB|10.687|13.477\n|TIGON|6.140|6.973\n|DeepRUOT|6.052|6.757\n|CytoBridge(Ours)|**6.013**|**6.644**\n\n> W2: Training Complexity\n- Thank you. We fully agree that directly optimizing four neural networks could be challenging. To address this, here we adopted the **two-phase pre-training strategy** detailed in Algorithm 2. This strategy effectively **decouples** the optimization problem: We first initialize the velocity ($v_\\theta$) and growth ($g_\\theta$) networks and the interaction potential ($\\Psi_\\theta$), providing a robust starting point for stable end-to-end training. This significantly reduces sensitivity to challenges associated with PINN losses. **We will expand the discussion of this strategy in the manuscript.**\n- Regarding hyperparameter selection, we found that lowering the mass loss weighting during pre-training improves distribution matching and performance. Additionally, extending training epochs for larger datasets enhances results. As detailed in Appendix C.2 and summarized in Table 14, we maintained most key hyperparameters consistent across diverse datasets, from synthetic to real-world scRNA-seq data, except for the scale parameter ($d_\\text{cutoff}$), which can be computed directly from each dataset. We also found that increasing the number of training epochs has no adverse effect, allowing users to use a default of (500, 100, 500) epoches. **We will include the more detailed discussion of the hyperparameter selection in the paper.**\n- To support this, **we successfully applied these hyperparameters** (Epoches {500, 100, 500}, with other parameters same in the table) to two additional datasets—a **100D EB dataset and a zebrafish spatiotemporal  (ST) dataset**—with results wil be presented below.\n\n> Q1 : Evaluate the integral during training?\n- Yes, the integral is evaluated during each training step. This is performed numerically using a standard ODE solver. Computed over ODE paths, this step incurs minimal overhead, ensuring computational efficiency without impeding the training process.\n\n>  Q2 : Interpretation of the Cell Interaction\n- Thank you. We fully acknowledge the significance of providing a clear biological interpretation of the cell-cell interactions modeled by CytoBridge. **Below, we present the biological interpretability of CytoBridge from three key perspectives:**\n  - **Detection of Cell-Cell Interactions Across Datasets:** CytoBridge’s interaction term enables the identification of the presence or absence of cell-cell interactions in a dataset. In our study, we analyzed four biological datasets: mouse hematopoiesis, EB, pancreatic differentiation, and EMT datasets. The first three datasets involve organoid-like systems cultured in vitro, where cell-cell interactions are plausible due to their three-dimensional, multicellular organization mimicking tissue-like environments. Conversely, the EMT dataset, derived from an in vitro cell line experiment, consists of isolated cells with minimal physical or signaling interactions. **Our algorithm successfully identified weak interaction patterns in the EMT dataset** (Appendix B.5), aligning with biological expectations, while **detecting stronger interactions in the organoid-like datasets**, consistent with their multicellular nature.\n  - **Interpretable Analysis of Interaction Contributions:** To enhance interpretability, we identified top 200 genes influenced most significantly by cell-cell interactions. Subsequent enrichment analysis of these genes revealed key pathways associated with the biological processes. Applying this to the mouse hematopoiesis dataset, we identified pathways related to positive regulation of leukocyte activation and cell-cell adhesion which are closely linked to hematopoiesis, interactions, and differentiation. **This demonstrates CytoBridge’s ability to uncover biologically relevant mechanisms.**\n  - **Extension to Spatiotemporal Transcriptomics**: Our framework is readily extensible to ST data. We applied CytoBridge to a zebrafish spatiotemporal ST dataset (*Dev. Cell*, 2022, PMID: 35512701). Further details are provided in our response to Reviewer 6tJf.\n- So the ability to detect interaction patterns, identify key contributing genes, and uncover relevant biological pathways highlights its significant potential. **We will expand the discussion of these points in the revised manuscript to ensure clarity.**\n\n| Pathway (Mouse Data) | p.adjust | Count |\n|:---|---:|---:|\n| regulation of T cell activation | 3.38e-12 | 23 |\n| positive regulation of cell activation | 1.42e-11 | 23 |\n| positive regulation of cell-cell adhesion | 5.81e-11 | 20 |\n| positive regulation of leukocyte activation | 1.9e-10 | 21 |\n| leukocyte cell-cell adhesion |3.18e-10 | 21 |\n| lymphocyte differentiation | 1.04e-09 | 21 |\n    \n> CytoBridge scale to higher dimensions?\n- Thank you. **Our method maintains scalability comparable to these existing neural ODE-based frameworks.**  In TrajectoryNet, MIOflow, Neural LSB and DeepRUOT, these methods have been tested up to 100 dimensions.\n-  **Following your suggestion, we have now evaluated CytoBridge on the EB data, extending the input from 50 to 100 PCs**. The results ($\\mathcal{W}_1$, time), reported below, demonstrate that our algorithm remains effective at this dimensionality. \n\n|Methods (100D EB)|t=1|t=2|t=3|t=4|Time|Memory|\n|-------|---|----|----|----|----|------|\n|MMFM|13.150|14.135|14.441|14.907|4min38s|0.6G|\n|MetaFM|11.699|13.398|14.037|14.727|5min50s|2.2G|\n|uAM|13.628|18.315|20.309|22.973|38s|1.5G|\n|SF2M|11.333|12.982|13.718|14.945|4min43s|0.7G|\n|MIOflow|11.387|12.331|11.905|12.908|6min19s|0.7G|\n|TIGON|10.547|12.926|13.897|14.535|7min33s|0.7G|\n|DeepRUOT|10.226|11.110|11.544|12.424|10min17s|3.7G|\n|CytoBridge|**10.217**|**11.070**|**11.505**|**12.368**|11min26s|6.3G|\n\n- This finding provides empirical evidence for the robustness of our method in scaling to more complex, higher-dimensional scenarios. However, at dimensions around 1000, our approach encounters challenges similar to those faced by other methods, highlighting the need for simulation-free training strategies, such as flow matching, to effectively handle ultra-high-dimensional settings.\n\n> Choice of p in RBM\n- Thank you. **We have conducted an ablation study to analyze its impact. We report the $\\mathcal{W}_1$ distances as below and will include the full results in the revised Appendix.** Generally, increasing $p$ slightly improves performance on evaluation metrics by providing a more accurate approximation of the interaction term. However, gains diminish at higher $p$ values, while computational costs increase. Thus, we selected $p=16$ as the default, balancing robust performance with computational efficiency.\n\n|p (Simulation Data)|t=1|t=2|t=3|t=4|\n|-|---|---|---|---|\n|2|0.025|0.033|0.048|0.094|\n|4|0.017|0.023|0.025|0.043|\n|8|0.016|0.019|0.022|0.035|\n|16|0.015|0.014|0.018|0.038|\n|32|0.015|0.016|0.020|0.030|\n|64|0.014|0.014|0.024|0.030|\n\n>  RBM inference vs standard inference?\n- Thank you. **We have conducted an ablation study to assess the RBM impact on model performance.** Results show that the full model (without RBM) and the RBM model ($p=16$) yield comparable performance, with the full model slightly better at certain time points. However, RBM significantly enhances computational efficiency, reducing the interaction term’s complexity from $\\mathcal{O}(N^2)$ (e.g., 10^6 for 1000 particles) to $\\mathcal{O}(pN)$. The tables demonstrate substantial reductions in memory and inference time with RBM. Given its comparable accuracy and scalability for large-scale single-cell datasets, RBM is a justified and essential component of our framework. Results will be detailed in the revised Appendix C.2.\n\n|Method (Simulation Data)|t=1|t=2|t=3|t=4|Time(s)|Memery(GB)|\n|-|---|---|---|---|---|---|\n|w/o RBM|0.015|0.013|0.023|0.029|0.568|22.1|\n|p=16|0.015|0.014|0.018|0.038|0.115|1.7|\n\n> Def 3.2. Figure 1 is never referenced in the text.\n- Thank you. We apologize for the typographical error. We will thoroughly review the manuscript to address any additional errors.\n\n> Limitations\n- Thank you. We will include the limitations and future work in the main text.\n\nThank you again!"}, {"review_id": "e2QnPDtwJK", "reviewer": "Reviewer_6tJf", "summary": "The authors introduce a method that accounts for cellular interactions while learning single-cell dynamics using Schrödinger bridges. The paper relies on the Unbalanced Mean Field Schrödinger Bridge (UMFSB) formulation, which allows for an interaction potential term modelling communication events along the dynamics. For this purpose, the authors implement **CytoBridge**, a deep-learning-based method to approximately solve the UMFSB problem with interaction terms. First, the authors define an ODE-based reformulation of the stochastic problem based on the associated probability-flow ODE. Then, they propose a simulation-based training approach where they approximate different components of the UMFSB objective with neural networks and train them with different losses that reflect the separate components of the minimization problem. In their experiments, the authors showcase their model on synthetic and real data. In the former setting, they show that the model recovers pre-defined interactions while capturing the cellular dynamics. In the latter scenario, the authors demonstrate improved modelling of the cellular transition, mass variation and velocity visualization.", "questions": "**Random batch method:** How do you think the batch size influences the model behaviour and performance in the random batch approach for the interaction term?\n\n**Computation of the interaction term during training.** During training, do you compute the interaction terms based on the real particles or the ones generated by integrating the system from the previous step?\n\n**Initial weights.** Is it reasonable to start from an initial uniform weighing scheme of $\\frac{1}{N}$? I imagined that the weights represent, in a way, the concentration of a biological state at a certain time point. Could you envision strategies for a more sophisticated initialization?\n\n**Target weights.** What is the ground truth for the target weights mentioned in line 218?\n\n**Algorithm - line 5.** Does $M_t$ in the algorithm mean that, in this phase, you only optimize the local mass constraint? Moreover, I find the notations in lines 12 and 13 of the algorithm a bit confusing. Do you learn a coupling for any couple of subsequent time points?", "limitations": "Limitations are reported in Appendix E.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "## Strengths\n\nThe idea at the core of the paper has strong potential, in my opinion. Novel approaches in the field of scRNA-seq are increasingly relying on spatial strategies to uncover cellular communication events, raising awareness that biological processes do not happen in isolated systems. I think that a similar approach as proposed in the paper would be very interesting in the field of spatial biology or other settings where explicit cellular interactions are measured and impactful on the biological processes at the single-cell level. I also like the math in the paper, I find it elegant and well-presented. \n\n\n## Weaknesses \n\nAlthough there surely are positive aspects to the papers, some general concerns led me to choose a provisional rejection score. I am looking forward to hearing more from the authors during the rebuttals. \n\n**Biological significance.** Although the paper has a dense methodological component, the application setting is the ultimate goal of the framework. I do agree that cellular interactions are crucial in understanding the system's dynamics, but here, in my opinion, the way communication is formulated is a bit ill-posed. If the model were applied in space, where communication has a strong impact and can be related to physical proximity, then I would find a lot of value in the formulation. However, since we are talking about dissociated data and cell culture, I do not see what the biological meaning of the interactions is.\n\n**Interactions in the experiments.** This is somewhat reflected in the choice of the experiments. The authors embed attractive, Lennard-Jones and repulsive potentials in the simulated data, but these are very idealized systems, and the effectiveness of the potentials in dissociated single-cell biology lacks a convincing justification to me. \n\n**The significance of the RBF kernel.** Along the lines of what I exposed above, I struggle to understand the choice of using an RBF function to enforce interactions between similar cells biologically. I guess the cell's dynamics across time can be influenced by different cell types, whose contribution would be penalized by the current choice for the potential. \n\n**Line 118.** Here, I would define $\\sigma$ and $\\mathbf{W}$ in a short sentence.\n\n**Line 122.** Where is $\\phi$ used on the definition? I would personally not include a specification about it if it does not appear in formulas. One could confuse it with the interaction potential. \n\n**Lines 137-139.** In my opinion, this sentence should be reformulated due to a lack of clarity and unclear punctuation. \n\n**Theorem 4.1.** For easier understanding, I would define $\\alpha$. In the current state of the paper, it is not introduced. \n\n**Lines 149.** Here, there is a typo. Trackable --> Tractable\n\n**Minor: citation style.** I recommend using `\\citet` for in-line citations to avoid the parentheses (e.g., line 82).\n\n**Algorithm stability and complexity.** While I really like the theoretical framework proposed by the authors, I feel that the optimization strategy proposed in the paper is not very stable. The approach requires a lot of heuristics in terms of pre-training components, approximations and loss balancing. I believe that, in a research setting where tractable and stable simulation-free models are being developed consistently, the methodological contribution loses a bit of relevance. One aspect that I would certainly include in the current paper formulation is some notion of run time and complexity compared to the competing algorithms. \n\n**Dependence on CFM.** Along the lines of the previous comment, I find the need to use CFM to pre-train the score sub-optimal, as you basically have to run a whole optimization round using a competing model to train CytoBridge properly. I would, at least, consider including CFM-based baselines using unbalanced OT, like [1].\n\n**Evaluation metrics.** I think it would be great if you referenced Appendix C.1 somewhere in the text. I was not familiar with the TMV metric and this broke the flow of the paper a bit for me. \n\n**Result presentation.** In my opinion, the result presentation could be improved. In Fig. 2, the legend says that crosses are the ground truth and dots are predictions. But I cannot see any crosses in the plots. Fig.1f has unlabelled axes, as well as all plots in Fig. 3. \n\n**Line 261.** I think the reference to the figure is swapped between potential and growth rate. \n\n**Waddington landscape.** I think the statement regarding the Waddington landscape could be better explained and presented. As of now, the authors simply point out that the model captures the right developmental potential, but how do I see this from the figure? Labelling the plot's axes would also help. \n\n**Baseline citation.** I think the MMFM baseline in the tables refers to the Multi-Marginal Flow Matching algorithm by (Rohbeck et al), so the citation necessitates a correction. \n\n**Description of the plots.** In general, I recommend that the authors always provide some qualitative understanding of why a plot shows the model is working well. For example, in line 276, the authors comment that the model learns growth rates properly. But how do I understand it from the plot in Figure 2b? What cell types are assigned a high growth rate, and does it make sense biologically? Also, the scores and interactions are hard to interpret if presented as stream plots on a UMAP. Maybe one could choose a more intuitive way to portray sinks and sources, or the growth direction?\n\nI hope I have not misinterpreted any crucial aspects of the paper. I am open to revising my assessment should the authors provide a convincing rebuttal.\n\n[1] Eyring, L., Klein, D., Uscidda, T., Palla, G., Kilbertus, N., Akata, Z. and Theis, F., 2023. Unbalancedness in neural monge maps improves unpaired domain translation. arXiv preprint arXiv:2311.15100.", "was_revised": true, "final_justification": "Dear AC, \n\nThank you for taking the time to handle this paper. After reading the other reviews, the rebuttals and the authors' responses to my concerns, I decided to increase my initial 2 to a 3. \n\nIn general, the method trained by the authors appears to work well, and the rebuttal reinforced this evidence via comparisons with the unbalanced OT-CFM model. \n\nHowever, I remain unconvinced by some aspects of the current version of the paper, which I will briefly list below:\n* I still believe that interactions in dissociated single-cell RNA-seq data are underdefined unless complemented with significant priors on the biology of the data. In other words, I do not think that the current model can identify realistic interactions dictated by molecular exchange and proximity, unless relying on some information on ligand-receptor pairs. \n* I still believe the algorithm is very convoluted and requires a preliminary optimization of an OT-CFM, a competing method. \n* The model explores very simplistic potentials (like LJ or standard attraction/repulsion) that are biologically unmotivated. I believe a more flexible definition of interactions is required. \n* In my opinion, the presentation of some of the figures requires an improvement (e.g., defining biological regions on the potential plots and replacing streamplots with more informative quantitative analysis). \n\nAll in all, in my opinion the quality of the paper would benefit a lot from another iteration of revisions, where the substantial feedback received by the reviewers is incorporated. \n\nThank you again for all your work in coordinating this process.", "rebuttal": "> Summary and Strengths\n- Thank you. We sincerely appreciate your recognition of the strong potential of our method. We address the challenge of **jointly** modeling cell-cell interactions and unbalanced stochastic biological effects using multiple single-cell snapshots data. This makes our work a new method in the integration of **interpretable cellular interactions** with **complex cell dynamics** within the optimal transport framework, offering both **great theoretical and practical potential** for exploring the **spatiotemporal dynamics and interactions of single cells**. \n\n> Weaknesses 1 and 2: Biological significance and Interactions\n- Thank you for your insightful comments. We fully agree that applying our CytoBridge framework to spatiotemporal transcriptomics data holds strong potential, and we appreciate the opportunity to address this point. Below, we outline our perspective on this application and its challenges, and explain our focus on scRNA-seq data analysis in the current work:\n  - **Challenges in ST:** Extending CytoBridge to ST data introduces complexities, particularly in **aligning** temporal and spatial scales across time-series slices due to technical and biological variability. Non-rigid or non-linear spatial alignment are needed and pose great difficulties.\n  - **Limited Methods for Joint Modeling:** To our knowledge, **very few studies** simultaneously model **spatiotemporal batch effects**, cell-cell interactions and complex dynamics (e.g., growth and stochastic) in this context, making the application of our framework to spatial transcriptomics directly very challenging. Applying CytoBridge to scRNA-seq data simplifies complexity, prioritizing cellular interaction modeling. We focus on this in the current work, designating detailed exploration of spatiotemporal data as future work.\n- Meanwhile, to show CytoBridge’s potential, we are very happy to present an example applying it to a **zebrafish ST dataset (Please see below)** (*Dev. Cell*, 2022, PMID: 35512701). And we believe CytoBridge hold strong biological significance from three key perspectives, even in the non-spatial data:\n  - **Detection of Cell-Cell Interactions Across Datasets:** Even in non-spatial dataset, CytoBridge’s ability to model interaction term enables the identification of the presence or absence of cell-cell interactions in a dataset. In our study, we analyzed four biological datasets: mouse hematopoiesis, EB, pancreatic differentiation, and EMT datasets. The first three datasets involve organoid-like systems cultured in vitro, where cell-cell interactions are plausible due to their three-dimensional, multicellular organization mimicking tissue-like environments. Conversely, the EMT dataset, derived from an in vitro cell line experiment, consists of isolated cells with minimal physical or signaling interactions. **Our algorithm successfully identified weak interaction patterns in the EMT dataset** (Appendix B.5), aligning with biological expectations, while **detecting stronger interactions in the organoid-like datasets**, consistent with their multicellular nature.\n  - **Interpretable Analysis of Interaction Contributions:** To enhance interpretability, we identified top 200 genes influenced most significantly by cell-cell interactions. Subsequent enrichment analysis of these genes revealed key pathways associated with the biological processes. On mouse hematopoiesis dataset, we identified pathways closely linked to hematopoiesis, interactions, and differentiation (Please see the reply to Reviewer idfK). **This demonstrates CytoBridge’s ability to uncover biologically relevant mechanisms.**\n  - **Extension to Spatiotemporal Transcriptomics**: We applied CytoBridge to a zebrafish spatiotemporal transcriptomics dataset, using 5.25 hpf and 10 hpf as input. We first compared the $\\mathcal{W}_1$ distances in the physical space (denoted as 'Space') and gene expression space (denoted as 'Gene') between CytoBridge and other methods. CytoBridge achieves better performance over other methods. Through interpretable analysis, we identified pathways related to somite development which are critical to zebrafish embryonic development. These findings align with known biological processes, showing the framework’s applicability to spatially resolved data.\n    \n|Methods (ST data)|Space|Gene|\n|-------|-----|----|\n|MMFM|0.247|5.208|\n|SF2M|0.265|5.423|\n|MetaFM|0.268|5.413|\n|uAM|0.177|6.499|\n|UOT-FM|0.227|5.173|\n|MetricFM|0.273|5.366|\n|MIOflow|0.263|4.720|\n|Tigon|0.352|4.979|\n|DeepRUOT|0.261|4.745|\n|Ours|**0.035**|**4.712**|\n\n|Pathway|p.adjust|Count|\n|---|---|---|\n|nucleolus|5.76e-05|12|\n|somite development|1.82e-04|10|\n|lipid transport|1.82e-04|12|\n|gastrulation|4.15e-04|12|\n|somitogenesis|4.15e-04|8|\n\n- So while more systematic ST dataset application requires further development especially in spatiotemporal alignment, **CytoBridge’s preliminary success on spatial data and its robust biological interpretability across diverse datasets highlight its potential.** We believe it will inspire important follow up research in the spatiotemporal dynamics research. **We will clarify these points in the revised manuscript.**\n    \n> The significance of the RBF kernel. \n- Thank you. The RBF kernel transforms the one-dimensional cell distance $d_{ij}$ into a high-dimensional feature vector. This allows our model to learn interactions between different cells by encoding interaction distances across multiple scales, rather than enforcing interactions only between similar cells. **We will add the detailed discussion in the revised version**.\n\n> Typo, citation, reference, presentation, evaluation metrics. \n- Thank you. We apologize for the typographical error and for confusion. **We will thoroughly review the manuscript to address any additional errors. We will update all related presentation to ensure clarity.**\n\n> Algorithm stability and complexity.\n\n- Thank you. We fully agree that flow matching is a very important area to expore, to our knowledge, simulation free jointly learning growth, interaction and velocity from snapshots data is challenging but **can be a very intersting future direction.** CytoBridge is **robust and does not require extensive, dataset-specific tuning**, enhancing the method's practical applicability. We have provided a detailed response regarding this to Reviewer rSSy. Please refer to there. Thank you for your understanding :).\n\n> Including UOT-FM\n\n- Thank you. The UOT-FM is an outstanding and very important work. We have now benchmarked UOT-FM. We show the $\\mathcal{W}_1$ distance on the mouse dataset as below. And also on 100D EB dataset (Please see our reply to Reviewer rSSy). **We will add the results of UOT-FM to all relevant tables in the paper.** The results show that CytoBridge is better on the key datasets.\n\n| Model | t=1 |t=2|\n| :--- | :--- |:---|\n| UOT-FM | 8.114 | 9.170 |\n| CytoBridge (Ours) | **6.013** | **6.644** |\n\n> Waddington landscape. \n- Thank you. The low-lying regions on this landscape correspond to areas of high probability density. These regions represent the stable states or attractors of the system. We apologize for the missing axis labels. We will clearly label the axes to make the visualization easier to understand.\n\n\n> Description of the plots.\n- Thank you. In the mouse dataset, the regions with high learned growth rates correspond to the hematopoietic stem cell (undifferentiated cell type) populations. **This is biologically consistent with the lineage tracing barcode results** (Nat. Mach. Intell., 2023, PMID: 38274364). **We will add more discussion in the revised version**.\n- We agree that stream plots can be challenging to interpret. Regarding the visualization of scores, we will highlight the cell types with higher stability. Additionally, we will highlight Figure 9’s analysis in the main text, which correlates the learned interaction force with cell transition velocity, directly illustrating whether interactions promote or inhibit differentiation, providing a more intuitive biological interpretation.\n\n> Questions: Random batch method\n- Thank you. **We have conducted an ablation study to analyze its impact, and we will include the full results in the revised Appendix.** Generally, increasing $p$ slightly improves performance on evaluation metrics by providing a more accurate approximation of the interaction term. However, gains diminish at higher $p$ values, while computational costs increase. Thus, we selected $p=16$ as the default, balancing robust performance with computational efficiency (Please see our reply to Reviewer idfk, thanks:)).\n\n> Computation of the interaction\n- Thank you. Forces need be calculated at each intermediate time step of the ODE solver, using the current, simulated particle positions. Relying only on fixed, real data points can not capture the dynamic influence of these forces between observed time points.\n\n> Initial weights. \n- Thank you. Mathematically, we use the empirical measure $\\mu_0^N=\\frac{1}{N}\\sum_{i=1}^N\\delta_{x_i}$ to approximate the true distribution $\\mu_0$, hence the uniform weights. This convergence to the true distribution is guaranteed when $N \\rightarrow \\infty$. We'll add this explanation to the paper. Thank you for pointing this out!\n\n> Target weights. \n- Thank you. As detailed in Appendix A.4, this weight, derived from the number of closest real data points, encourages the growth network to assign higher weights to particles moving into denser state space regions. **We'll elaborate more on the point in the revised version.**\n\n> Algorithm - line 5. \n- Yes, correct! In pre-training, we optimize the local mass constraint ($M_t$) for more flexible, fine-grained guidance. Line 12's notation means we sample $(x_0, x_1)$ pairs from the optimal transport plan between consecutive time distributions to build Brownian bridges for score matching. **We'll add the explanations in the revised version.**\n\nThank you again!"}, {"review_id": "OlRG71QEME", "reviewer": "Reviewer_SjPe", "summary": "The paper proposes a new mathematical framework—Unbalanced Mean-Field Schrödinger Bridge (UMFSB)—for learning stochastic, unbalanced cell-state dynamics with explicit cell–cell interactions from snapshot data such as scRNA-seq. The authors present deep-learning method is suggested, CytoBridge, which provides an approximated solution to the mathematical formulation.  The model is compared to state-of-the art methods, with experiments over (i) synthetic gene-regulatory examples with several interaction types, and (ii) four real scRNA-seq time-series (mouse hematopoiesis, embryoid body, pancreatic β-cell differentiation, and TGFB1-induced EMT).", "questions": "Following the presented weaknesses can the authors:\n1. provide additional benchmarks highlighting the significance of cytoTrace.\n2. extend the biological analysis to support interpretability claims.", "limitations": "Yes", "rating": 3, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 1, "originality": 2, "strengths_and_weaknesses": "*Quality*:\n\n_strengths_—\n- Solid formulation: UMFSB clearly nests regularised OT, mean-field SB, RUOT and dynamical OT as special cases.\n- Extensive experiments and evaluations: 5 datasets, ablations, and hold-out tests.\n\n_weaknesses_—  \n- Practical quality of the presented experiments is limited-quantitatively showcasing minimal improvement. \n- Thought framed around biological contribution the biological story is missing (from motivation to practical impact).\n\n*Clarity*:\n\n_strengths_— \n- The paper is mostly well-structured: preliminaries → theory → algorithm → experiments.\n- All details regarding implementation and evaluations are provided, \n\n_weaknesses_—  \n- Lengthy exposition makes it hard to follow and grasp contribution.\n- Experiments/results section is short relative to mathematical parts, with biological interpretation sections completely missing, failing to provide justification for the contribution claims.\n\n*Significance*:\n\n_strengths_— \n- Addresses a gap in transport based single-cell trajectory inference methods which do not account explicitly for cell–cell interactions\n- Conceptually, CytoBridge provides interpretable growth and interaction functions, which can be used for biological insight.\n\n_weaknesses_—practical significance is very unclear\n- Quantitative results show minimal improvement over compared methods. \n- Comparison neglects standard single-cell trajectory inference methods, rna velocity (though scVelo is used for visualizatin), and methods purposed in original publication. Further, it will be valuable to compare the method and relate to TICCI (\"Trajectory Inference with Cell–Cell Interactions\", Fu et al., 2025)\n- The provided visualizations of learnt velocity, growth rates, and interactions over some lower dimensional representation (differing between modalities) do not provide direct biological insight. To highlight their relevance one could evaluate these quantities with respect to the known biology (as all datasets were previously extensively studied). \n\n*Originality*:\n\n_strengths_—introduces a genuinely new mathematical framework—Unbalanced Mean-Field Schrödinger Bridge (UMFSB)—the first time unification of mass imbalance, stochasticity, and particle interactions, applicable for studying cellular dynamics.  \n\n_weaknesses_—the application to single-cell field joins an extensive line of tools and failing to account for the scale and complexity of current datasets.", "was_revised": true, "final_justification": "I acknowledge the detail and effort put in the rebuttal. On my end am afraid that the authors have not addressed the weaknesses I previously highlighted and make statements that I believe are not supported by the presented results.\n\nTo clarify, I have never suggested that \"cell-cell interactions across datasets are irrelevant.\" Instead, I have stated again that simply estimating the magnitude of interaction is not enough to claim novel biological insight and requested further explanation on the \"biological interpretation\". Generally, the examples given only show broad trends and do not demonstrate an understanding of lineage-specific outcomes.\n\nIn summary, while I agree with the authors that the method shows significant promise, the current results do not yet support several of the claims made.\n\nThank you, and best wishes for the next steps.", "rebuttal": "> Summary and Strenghts\n\n- Thank you. We sincerely appreciate your recognition of the strengths in our paper. \n-  scRNA-seq offers only **a static snapshot** of gene expression, lacking the temporal information needed to understand cell state transitions. Our work, similar to recent advances like WaddingtonOT (*Cell, 2019*) and Moscot (*Nature, 2025*), focuses on **temporally resolved scRNA-seq** to capture gene expression profiles **across multiple time points**.\n- We address the challenge of **jointly** modeling cell-cell interactions and unbalanced stochastic biological effects using multiple single-cell snapshots data. This is an **early and essential** step in the integration of **interpretable cellular interactions** with **complex cell dynamics** within the **optimal transport** framework, offering both **great theoretical and practical potential** for exploring the **spatiotemporal dynamics and interactions of single cells**. \n\n> Compared to Works in Trajectory Inference\n- Thank you for your valuable feedback. **We are pleased to clarify the distinctions between various methods in trajectory inference and their applicable contexts.** In single-cell transcriptomics, methods vary depending on the data type:\n- **Single Time-Point Snapshot Data:** In this context, trajectory inference methods are broadly categorized into two groups:\n  - **Pseudotime-Based Methods** (e.g., Monocle, Slingshot, PAGA, Diffusion Map) infer trajectories by ordering cells along a pseudotime axis but often **require specifying a starting point**, limiting their flexibility.\n  - **RNA Velocity-Based Methods** (e.g., scVelo, VeloAE, VeloVI, DeepVelo) leverage **spliced and unspliced** counts to estimate dynamics but are constrained by the need for such data. **Both approaches infer dynamics from existing data without generating new cell states.**\n- **Multi-Time-Point Snapshot Data**: Our work addresses a distinct scenario involving single-cell sequencing data **across multiple time points**. Due to the destructive nature of technology, this can be framed as a **generative modeling** task, where dynamics are inferred from distributions at different time points and could be interpolated at unseen time points. Once learned, these dynamics enable the generation of intermediate cell states. Methods like **optimal transport**, which have gained significant attention in computational systems biology and machine learning, are well-suited for this task.\n- **Evaluation and Comparability:** Our generative approach, unlike pseudotime or RNA velocity methods, evaluates performance using metrics for **generated distributions, which traditional methods cannot produce**. This fundamental difference in data and modeling objectives, along with our **hold-out experiments**, renders direct comparisons with traditional trajectory inference methods infeasible. **Very few studies**, to our knowledge, compare these method types in the same context. A recent review (PMID: 40422408) discusses modeling approaches for different data types. The table below summarizes these distinctions:\n\n| Method | Data | Key | Generative | Evaluation|\n|-------------|-----------|---------------------|-----------------------|--------------------|\n| Pseudotime-Based| Single snapshot or merged snapshots without using time labels| Infers trajectories by ordering cells; requires starting point | No | Pseudotime accuracy with true time labels |\n| RNA Velocity-Based| Single snapshots with spliced/unspliced counts | Estimates dynamics using RNA splicing kinetics | No | Velocity consistency |\n| Optimal Transport-Based (Ours) | Multi-time-point snapshots | Infers dynamics and generates new cell states | Yes | Generated distribution metrics |\n\n- We greatly appreciate your suggestion to consider the TICCI method. TICCI is a highly significant and outstanding work. Moreover, TICCI uses CellChat to compute intercellular communication patterns, and this also can be integrated into CytoBridge’s function $k$, enhancing its compatibility with biological priors and practical value. Although direct comparisons are inappropriate due to differing objectives and evaluation metrics, **we are very happy to incorporate TICCI as a reference in our study.** We have applied TICCI to our datasets to validate our biological intepretation and will **discuss it as a relevant work**  in the revised manuscript. **Additionally, we will expand the discussion to provide a more thorough comparison between these methods.** \n\n> Limitated improvement, Biological Significance and Interpretation\n- Thank you. We address your comments as follows:\n- **Comprehensive Benchmarking:** In our study, we compared CytoBridge against **nine representative methods (including two newly added methods)**, including neural ODE-based, flow matching-based, multi-marginal, and unbalanced approaches. To our knowledge, **these encompass the state-of-the-art in the field.** CytoBridge demonstrates **clear performance improvements** across multiple datasets, underscoring its effectiveness. \n- **Novel Contribution:** And we believe our work is the very first to provide a framework that simultaneously learns explicit cell-cell interaction dynamics alongside unbalanced and stochastic effects, a critical biological process difficult be modeled by existing methods. This enables direct inference of cell-cell interactions from snapshot data, **offering novel biological insights that constitute the primary significance of our contribution.**\n- **Strong Potential:** As noted by Reviewer 6tJf, CytoBridge holds substantial potential for spatiotemporal transcriptomics. **We have incorporated an additional spatial transcriptomics dataset and conducted further analyses**, demonstrating CytoBridge’s strong biological relevance across both spatial and non-spatial data from three key perspectives:\n  - **Detection of Cell-Cell Interactions Across Datasets:** CytoBridge’s interaction term enables the identification of the presence or absence of cell-cell interactions in a dataset. In our study, we analyzed four biological datasets: mouse hematopoiesis, EB, pancreatic differentiation, and EMT datasets. The first three datasets involve organoid-like systems cultured in vitro, where cell-cell interactions are plausible due to their three-dimensional, multicellular organization mimicking tissue-like environments. Conversely, the EMT dataset, derived from an in vitro cell line experiment, consists of isolated cells with minimal physical or signaling interactions. **Our algorithm successfully identified weak interaction patterns in the EMT dataset** (Appendix B.5), aligning with biological expectations, while **detecting stronger interactions in the organoid-like datasets**, consistent with their multicellular nature.\n  - **Interpretable Analysis of Different Contributions:** To enhance interpretability, we identified top 200 genes influenced most significantly by cell-cell interactions. Subsequent enrichment analysis of these genes revealed key pathways associated with the biological processes. Applying this to the mouse hematopoiesis dataset, we identified pathways closely linked to hematopoiesis, interactions, and differentiation (Please see our reply to Reviewer idfk). Moreover, we calculated the gradients of growth network with respect to genes to identify key genes that contribute to growth dynamics, including **Meis1, Nfkb1, and Rap1b**. Meis1 regulates hematopoietic stem cell proliferation and self-renewal, Nfkb1 promotes cell cycle entry and survival as a signaling hub, and Rap1b drives cell division via pathways like MAPK, consistent with established biological knowledge. **This demonstrates CytoBridge’s ability to uncover biologically relevant mechanisms.** \n  - **Extension to Spatiotemporal Transcriptomics**: We applied CytoBridge to a zebrafish spatiotemporal transcriptomics dataset (*Dev. Cell*, 2022, PMID: 35512701), using 5.25 hpf and 10 hpf as input. We find that CytoBridge is better over other methods. Through interpretable analysis, we identified pathways related to somite development which are critical to zebrafish embryonic development (Please see our reply to Reviewer 6tJf). These findings align with known biological processes, showing the framework’s applicability to spatially resolved data.\n- So, **CytoBridge’s success on spatial data and its robust biological interpretability across diverse datasets highlight its potential.** We believe it will inspire important follow up research in the spatiotemporal dynamics research. **We will clarify these points in the revised manuscript.**\n\n\n> Biological Motivation and Practical Impact.\n\n- Thank you. As outlined in the **related works section** of our manuscript, simultaneously modeling cell-cell interactions and complex dynamics (e.g., unbalanced stochastic effects) within the optimal transport framework is well-motivated for temporally resolved snapshot data. We have demonstrated the effectiveness of our algorithm across five biological datasets,including one spatiotemporal datasets, highlighting its practical impact and biological relevance. \n\n> scale and complexity\n- Thank you. We clarify that CytoBridge was evaluated on large-scale, real-world datasets, such as the pancreas dataset with **51,274 cells** across **8 time points**, and also a **spatiotemoral transcriptomics** achieving state-of-the-art performance. To address complexity, we further added experiments with **100 dimensional space**, which is **comparable** to the most of methods in this domain, confirming robust performance in high-dimensional settings (results included in the revised manuscript). We **comprehensively evaluated computational time and memory** usage, confirming they are acceptable (Please see our response to Reviewer rSSy).\n- Thus, we believe CytoBridge is both **methodologically innovative** and capable of **handling the scale and complexity** of current single-cell datasets.\n\nThank you again!"}, {"review_id": "LDhXlAJVRq", "reviewer": "Reviewer_rSSy", "summary": "This paper introduces the **Unbalanced Mean-Field Schrödinger Bridge (UMFSB)** framework—an approach for modeling unbalanced, stochastic interaction dynamics from snapshot data, mainly applicable to single-cell trajectories.  \nThe authors present **CytoBridge**, a practical instantiation that approximates the UMFSB problem using four neural networks:   **v** – velocity field,  **g** – growth/decay rate,  **Φ** – interaction potential,   **s** – score (∇ log ρ).  These networks are trained jointly with a composite loss that combines an energy term, a reconstruction term, and a physics-informed Fokker-Planck residual.  Experiments on standard single-cell benchmarks demonstrate improved reconstruction accuracy over previous optimal-transport and Schrödinger-bridge baselines.", "questions": "1. **Hyper-parameter sweep** – How extensive was the sweep, and what procedure was used to select the final values? Would a similar search be prohibitive for higher-dimensional data, new datasets, or different settings?  \n2. **Missing baseline** – Could you compare against Metric Flow Matching[2], or explain why this method was not included?  \n3. **Computation** – Please provide approximate training times: wall-clock duration and peak GPU memory for the largest experiment.", "limitations": "Yes, section Limitations and Further Directions is on point.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "## Strengths  \n- **Addresses a relevant problem** – The method explicitly models cell–cell interactions and unbalanced data (cell growth/death), aspects that most existing approaches ignore.  \n- **Novel theoretical framework** – This is the first dynamic OT formulation that *simultaneously* supports stochastic regularisation (σ > 0), mean-field interactions (kernel *k*, potential Φ), and unbalanced mass creation/destruction (*g*).  \n- **Mathematical soundness** – The proof of Theorem 4.1 and the derivation of Proposition 5.1 are both correct.  \n- **Practical implementation** – Using four separate networks (*v*, *g*, *s*, Φ) is a reasonable design choice; it covers all the relevant phenomena, whereas previous methods implemented only subsets of these components.  \n- **Promising empirical performance** – The reported results look encouraging (see caveats below).  \n\n---\n\n## Weaknesses  \n- **Training complexity** – The method relies on Neural-ODE simulation, a two-phase pre-training schedule, and many hyper-parameters (e.g., loss weights λ, RBF bandwidth β, FP-penalty weight). This complexity makes the model difficult to train in higher dimensions and potentially sensitive to specific hyper-parameter choices; the ablation studies only partially address these concerns.  \n- **Training time** – The paper does not report wall-clock or GPU times. Long training cycles could limit practical adoption, especially given the emergence of simulation-free single-cell methods that train much faster (e.g., OT-CFM [1], OT-MFM [2]).  \n- **Omitted baselines** – While MIOFlow (Huguet et al., 2022) is included, newer simulation-free and manifold-aware baselines such as **Metric Flow Matching** [2] are missing. Because Metric Flow Matching is currently state of the art on the EB dataset, it should be discussed and compared.  \n\n---\n\n[1] Tong, A. *et al.* (2023). *Conditional Flow Matching: Simulation-Free Dynamic Optimal Transport*, arXiv:2302.00482.  \n[2] Kapusniak, K. *et al.* (2024). *Metric Flow Matching for Smooth Interpolations on the Data Manifold*, NeurIPS 37, 135011-135042.", "was_revised": true, "final_justification": "The rebuttal convincingly addresses my main concerns. The authors supply concrete training-time figures, show that settings transfers to larger datasets, and add the missing baselines. With these clarifications, the paper now offers a well-rounded and impactful contribution. I therefore raise my score to 5.", "rebuttal": "> Summary and Strengths\n- Thank you. We sincerely appreciate your recognition of the strengths in our paper. In this study, we address the challenge of **jointly** modeling cell-cell interactions and unbalanced stochastic biological effects using multiple single-cell snapshots data. This introduces a new framework in the integration of **interpretable cellular interactions** with **complex cell dynamics** within the optimal transport framework, offering both **great theoretical and practical potential** for exploring the **spatiotemporal dynamics and interactions of single cells**. \n\n> •\tTraining complexity and hyperparameter selection\n- Thank you. Here we adopted the **two-phase pre-training strategy** detailed in Algorithm 2. This strategy effectively **decouples** the optimization problem: We first initialize the velocity ($v_\\theta$) and growth ($g_\\theta$) networks and the interaction potential ($\\Psi_\\theta$), providing a robust starting point for stable end-to-end training. This significantly reduces sensitivity to challenges associated with PINN losses. **We will expand the discussion of this strategy in the manuscript.**\n- Regarding hyperparameter selection, we found that lowering the mass loss weighting during pre-training improves distribution matching and performance. Additionally, extending training epochs for larger datasets enhances results. As detailed in Appendix C.2 and summarized in Table 14, we maintained most key hyperparameters consistent across diverse datasets, from synthetic to real-world scRNA-seq data, except for the scale parameter ($d_\\text{cutoff}$), which is computed directly from each dataset. We also found that increasing the number of training epochs has no adverse effect, allowing us to use a default of (500, 100, 500) epoches. **We will include the more detailed discussion of the hyperparameter selection in the paper.**\n- To support this, **we successfully applied these hyperparameters** (Epoches {500, 100, 500}, with other parameters same in the table) to two additional datasets—a **100D EB dataset and a zebrafish spatiotemporal  (ST) dataset**—with results wil be presented below.\n- This demonstrates that CytoBridge **does not require extensive, dataset-specific tuning and that the presented hyperparameters serve as a reliable default configuration**, enhancing the method's practical applicability.\n\n> The paper does not report wall-clock or GPU times. Long training cycles could limit practical adoption, especially given the emergence of simulation-free single-cell methods that train much faster (e.g., OT-CFM [1], OT-MFM [2]).\n- Thank you. We fully agree that flow matching's simulation-free approach enables much faster computation. To our knowledge, jointly learning growth, interaction, and velocity from snapshot data in a simulation-free manner is challenging **but represents a highly promising future research direction.** Below, we present the computational time and memory usage for each method on the 100D EB dataset. Our results show that CytoBridge requires a training time of 11 minutes, comparable to other simulation-based methods, with a peak GPU memory usage of 6.3 GB. On the largest pancreatic $\\beta$-cell differentiation data with **51,274** cells across **8** time points, CytoBridge requires a training time of 23 minites, with a peak GPU memory of 3.1 GB.  Given the dataset's dimensionality and scale, we consider this computational cost to be acceptable. **We will list the training time and memory usage in Appendix C.**\n\n|Methods (100D EB Data)|Time|Memory|\n|-------|----|------|\n|UOT-FM|3min27s|0.6G|\n|MMFM|4min38s|0.6G|\n|MetaFM|5min50s|2.2G|\n|uAM|38s|1.5G|\n|SF2M|4min43s|0.7G|\n|MetricFM|2min16s|4.2G|\n|MIOflow|6min19s|0.7G|\n|TIGON|7min33s|0.7G|\n|DeepRUOT|10min17s|3.7G|\n|CytoBridge (ours)|11min26s|6.3G|\n\n\n\n> •\tOmitted baselines – While MIOFlow (Huguet et al., 2022) is included, newer simulation-free and manifold-aware baselines such as Metric Flow Matching [2] are missing. Because Metric Flow Matching is currently state of the art on the EB dataset, it should be discussed and compared.\n- Thank you for pointing this out. Metric Flow Matching is a very important and outstanding contribution. This is a critical comparison that will greatly strengthen our paper. We have now benchmarked CytoBridge against Metric Flow Matching and found that our method achieves improved performance on the key datasets. We show the $\\mathcal{W}_1$ distance on the 100D EB dataset as below, and **we will add the results of Metric Flow Matching to all relevant tables in the revised version.** \n- We also applied CytoBridge to a zebrafish **spatiotemporal** transcriptomics dataset (*Dev. Cell*, 2022, PMID: 35512701), using 5.25 hpf and 10 hpf as input. We first compared the $\\mathcal{W}_1$ distances in the physical space (denoted as 'Space') and gene expression space (denoted as 'Gene') between CytoBridge and other methods. As shown below, CytoBridge achieves improved performance over other methods. To enhance interpretability, we identified top 200 genes influenced most significantly by cell-cell interactions. Subsequent enrichment analysis of these genes revealed key pathways associated with the biological processes. Through interpretable analysis, we identified pathways related to somite development which are critical to zebrafish embryonic development. **These findings align with known biological processes, underscoring the framework’s applicability to spatially resolved data.** \n- We also applied CytoBridge to the mouse hematopoiesis dataset, identifying pathways associated with positive regulation of leukocyte activation and cell-cell adhesion, key processes in hematopoiesis, interactions, and differentiation. Moreover, we calculated the gradients of growth network with respect to genes to identify key genes that contribute to growth dynamics, including **Meis1, Nfkb1, and Rap1b**. Meis1 regulates hematopoietic stem cell proliferation and self-renewal, Nfkb1 promotes cell cycle entry and survival as a signaling hub, and Rap1b drives cell division via pathways like MAPK, consistent with established biological knowledge. **This demonstrates CytoBridge’s ability to uncover biologically relevant mechanisms.**\n- So, **CytoBridge’s preliminary success on spatial data and its robust biological interpretability across diverse datasets highlight its potential.** We believe it will inspire important follow up research in the **spatiotemporal** dynamics research. **We will clarify these points in the revised manuscript.**\n\n\n|Methods (100D EB Data)|t=1|t=2|t=3|t=4\n|-------|---|----|----|----\n|UOT-FM|10.757|12.799|13.761|15.657\n|MMFM|13.150|14.135|14.441|14.907\n|MetaFM|11.699|13.398|14.037|14.727\n|uAM|13.628|18.315|20.309|22.973|38s\n|SF2M|11.333|12.982|13.718|14.945\n|MetricFM|10.806|12.348|13.622|16.801\n|MIOflow|11.387|12.331|11.905|12.908\n|TIGON|10.547|12.926|13.897|14.535\n|DeepRUOT|10.226|11.110|11.544|12.424\n|CytoBridge (ours)|**10.217**|**11.070**|**11.505**|**12.368**\n\n\n|Methods (Zebrafish ST data)|Space|Gene|\n|-------|-----|----|\n|MMFM|0.247|5.208|\n|SF2M|0.265|5.423|\n|MetaFM|0.268|5.413|\n|uAM|0.177|6.499|\n|UOT-FM|0.227|5.173|\n|MetricFM|0.273|5.366|\n|MIOflow|0.263|4.720|\n|Tigon|0.352|4.979|\n|DeepRUOT|0.261|4.745|\n|Ours|**0.035**|**4.712**|\n\n|Pathway (Zebrafish ST data)|p.adjust|Count|\n|---|---|---|\n|nucleolus|5.76e-05|12|\n|somite development|1.82e-04|10|\n|lipid transport|1.82e-04|12|\n|gastrulation|4.15e-04|12|\n|somitogenesis|4.15e-04|8|\n\n| Pathway (Mouse Data) | p.adjust | Count |\n|:---|---:|---:|\n| regulation of T cell activation | 3.38e-12 | 23 |\n| positive regulation of cell activation | 1.42e-11 | 23 |\n| positive regulation of cell-cell adhesion | 5.81e-11 | 20 |\n| positive regulation of leukocyte activation | 1.9e-10 | 21 |\n| leukocyte cell-cell adhesion |3.18e-10 | 21 |\n| lymphocyte differentiation | 1.04e-09 | 21 |\n\n> Minor typo\n\n- Thank you. We apologize for the typo. **We will thoroughly review the manuscript to address any additional errors and will update all related presentation to ensure clarity.**\n\nThank you again!"}]

Author rebuttal:

Meta-review: This paper introduces CytoBridge, a framework for modeling single-cell dynamics that jointly accounts for unbalanced mass, stochasticity, and cell-cell interactions. This paper presents the Unbalanced Mean-Field Schrödinger Bridge (UMFSB), a new mathematical framework for inferring continuous-time cellular dynamics from discrete snapshots. The core innovation is its ability to simultaneously model three key aspects of biological systems: (1) stochasticity, (2) unbalanced populations, and (3) cell-cell interactions. To make this framework practical, the authors propose CytoBridge, a deep learning algorithm that approximates the UMFSB solution by jointly training four neural networks representing the system's velocity, growth, interaction potential, and score. The authors claim that CytoBridge can accurately reconstruct developmental landscapes and identify meaningful biological patterns, outperforming existing methods that ignore one or more of these critical components.

However, there are some unresolved core concerns. Despite significant effort, the rebuttal did not resolve the fundamental concerns of reviewers 6tJf and SjPe. They remained unconvinced by the biological significance of the learned interactions, arguing that the concept is ill-posed in dissociated data without stronger priors and that the validation did not provide novel insights. Reviewer SjPe explicitly stated, "the current results do not yet support several of the claims made." I agree with both these reviewers, the biological priors used and significance of learned interactions are simply not clear at this point in time. It seems very unclear that cell-cell interactions can be inferred without additional assumptions from dissociated single-cell RNA-seq data.

The discussion solidified a critical split in the reviews. Two reviewers were satisfied with the new experiments and raised their scores, the other two maintained borderline reject stances, unconvinced by the biological significance. While I agree that the biological significance is quite questionable, I believe the method has merit and is a step in the right direction for trajectory inference methods. For this reason, despite its faults and uncertainties I recommend acceptance of this work.
Decision: Accept (poster)