WFR-FM: Simulation-Free Dynamic Unbalanced Optimal Transport
Authors: Qiangwei Peng
Zihan Wang
Junda Ying
Yuhao Sun
Qing Nie
Lei Zhang
Tiejun Li
Peijie Zhou
Venue: ICLR 2026 Poster (2026)
Keywords: flow matching; unbalanced optimal transport; Wasserstein-Fisher-Rao
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: Solve the dynamic unbalanced optimal transport WFR problem by using flow matching

Abstract: The Wasserstein–Fisher–Rao (WFR) metric extends dynamic optimal transport (OT) by coupling displacement with change of mass, providing a principled geometry for modeling unbalanced snapshot dynamics. Existing WFR solvers, however, are often unstable, computationally expensive, and difficult to scale. Here we introduce \textbf{WFR Flow Matching (WFR-FM)}, a simulation-free training algorithm that unifies flow matching with dynamic unbalanced OT. Unlike classical flow matching which regresses only a transport vector field, WFR-FM simultaneously regresses a vector field for displacement and a scalar growth rate function for birth–death dynamics, yielding continuous flows under the WFR geometry. Theoretically, we show that minimizing the WFR-FM loss exactly recovers WFR geodesics. Empirically, WFR-FM yields more accurate and robust trajectory inference in single-cell biology, reconstructing consistent dynamics with proliferation and apoptosis, estimating time-varying growth fields, and  applying to generative dynamics under imbalanced data. It outperforms state-of-the-art baselines in efficiency, stability, and reconstruction accuracy. Overall, WFR-FM establishes a unified and efficient paradigm for learning dynamical systems from unbalanced snapshots, where not only states but also mass evolve over time. The Python code is available at <https://github.com/QiangweiPeng/WFR-FM>.

Published as a conference paper at ICLR 2026
WFR-FM: SIMULATION-FREE DYNAMIC UNBAL- ANCED OPTIMAL TRANSPORT
Qiangwei Peng' ; Zihan Wang Junda Ying"; Yuhao Sun Qing 6677, Lei Zhang?3At Tiejun Lil & Peijie 24.99
'LMAM and School of Mathematical Sciences, Peking University
Center for Quantitative Biology, Peking University
Beijing International Center for Mathematical Research, Peking University
*Center for Machine Learning Research, Peking University
NNSF-Simons Center for Multiscale Cell Fate Research, University of California, Irvine 'Department of Developmental and Cell Biology, University of California, Irvine 'Department of Mathematics, University of California, Irvine & National Engineering Laboratory for Big Data Analysis and Applications, Beijing
9 AI for Science Institute, Beijing
ABSTRACT
The Dasserstein-Fisher-Be (WFR) metric extends dynamic optimal transport (OT) by coupling displacement with change of mass, providing principled ge- ometry for modeling unbalanced snapshot dynamics. Existing WFR solvers, how- ever, are often unstable, computationally expensive, and difficult to scale. Here we introduce WFR Flow Matching (WFR-FM), simulation free training algo- rithm that unifies flow matching with dynamic unbalanced OT. Unlike classical flow matching which regresses only a transport vector field, WFR-FM simultane- ously regresses vector field for displacement and a scalar growth rate function for birth-death dynamics, yielding continuous flows under the WFR geometry. Theoretically, show that minimizing the WFR-FM loss exactly recovers WFR geodesics. Empirically. WFR-FM yields more accurate and robust trajectory infer- ence in single-cell biology, reconstructing consistent dynamics with proliferation and apoptosis, estimating time-var vvrying growth fields, and applying to generative dynamics under imbalanced data. It outperforms state-of-the-arr baselines in ef- ficiency, stability, and reconstruction accuracy. Overall, WFR-FM establishes a unified and efficient paradigm for learning dynamical systems from unbalanced snapshots, where not only states but also mass evolve over time. The Python code is available at
Peannngwwweeeinn///aeennn Pengee/i-//////T /NFR-FM
INTRODUCTION
Reconstructing continuous dynamics from limited observations is central challenge in many sci- entific domains (Chen et al., 2018). In single-cell transcriptomics, destructive sequencing protocols and experimental costs typically restrict measurements
small number of temporal snapshots (Zheng et al. [017]] Chen et al. 2022a).. The problem of trajectory inference is therefore to recover the underlying cellular dynamics from such sparse single-cell RNA sequencing (scRNA-seq) snap- shots (Schiebinger et al., 2019 Tong et al.) 2020). Optimal transport (OT) has then emerged
powerful framework for modeling single-cell population across time-series ScRNA-seq data (Bunne al. 2024| Zhang et al.) 2025b). Existing approaches
be broadly divided into two categories: methods based on static OT maps (Schiebinger et al. 2019 Klein ct al. 2025 Halmos et al. 2025). which align distributions at different time points without explicitly modeling intermediate dynam-
ics, and methods based on dynamic OT (Tong et al. 2020 Huguet et al.) 2022 Zhang et al. 2024a
Equual contribution
Cooresponding authors: pj pjzhou @acku.u@kkk... @pku. edu.cn
qnie@uci. edu,
zhang1 mangl1@mahh..mmm seaah.pkt. pku edu.cn, (tieli,

---

Published as a conference paper at ICLR 2026
Peng et al., 2024 Tong et al. 2023 Gu al. 2025 which reconstruct continuous flows between
snapshots. While dynamic OT methods provide richer dynamical insights, they are often realized via neural ordinary differential equations (ODEs) and continuous normalizing flows, leading to high computational cost and training instabilities due to repeated ODE simulations (Yan et al. 2019).
To accelerate learning of continuous normalizing flows, flow matching (FM) Lipmman et al. 2023
Liu et al. 2022 Pooladian et al.] 2023] Albergo & Ganden-Eijnden 2022) has been proposed as
simulation-free frree training paradigm. By directly regressing the drift field of an ODE, FM enables stable and scalable training without the need for integration. This motivated several works combin- ing OT with FM. which approximate dynamic OT simulation- free manner (Tong et al.) 2024a] Klein et al. 2024; Rohbeck et al., 2025] Eyring et al.] (024). Such approaches enable the efficient recovery of transport-basee dynamies, avoiding the prohibitive cost of repeated ODE solvers.
However, single-cell dynamics are not coss-conssrving: cells can proliferate and undergo apoptosis, resulting in unbalanced distributions across time (Sha et al., (0024). The Wasserstein- Feenn--Fhhhrreee (WFR) metric (Liero et al. [016,]
/ Chizat et al. [20 8aj Kondratyev et al., (016)) extends dynamic OT
by coupling transport with mass creation and annihilation, providing a principled geometry for such unbalanced dynamics (Sha et al. 2024 Peng al. 2024 Zhang et al., 2025a [un et al. 2025). This
has inspired several extensions of FM to unbalanced settings (Eyring et al. 2024; Cao et al. 2025]
Corso et al.
2025 Wang et al. 2025). However, existing approaches typically focus regressing
only velocity fields and neglecting explicit modeling of growth dynamics (Eyring et al.] 2024 Corso
ett al., 2025 Cao al. 2025). The recent VGFM (Wang et al. (025) introduces innovative flow matching framework that jointly learns both transport and growth. Its formulation separates birth-death dynamics from transport and then uses modified dynamics approximate simultancous evolution, which may depart from the geometry of dynamic unbalanced OT. In addition, while being effective, VGFM still relies on ODE simulation during post-training refinement and therefore is not fully simulation-fre frree Action Matching (AM) (Neklyudov et al. (023) proposes novel simulation-free approach that can handle unbalanced dynamics. However, it assumes the access to a continuous curve of densities, which might be challenging in scRNA-seq experiments. The Wasserstein Lagrangian Flows (WLF) framework (Neklyudov et al., 2024 addresses the issuc with an inner-outer optimization design, albeit the cost of introducing additional computational costs.
To address these challenges, we introduce WFR-FM, which combines the core ideas of unbalanced OT and FM solve dynamic unbalanced OT under the WFR geometry in simulation-free manner. WFR-FM jointly regresses both a transport vector field and a growth rate function, yielding continu- ous flows that explicitly model displacement and birth-death dynamics. We show that this approach produces more accurate WFR-OT flows than existing ODF-based solvers [Sha et al. (2024); Zhang flow matching methods Tong et al. 2024b); Kapusniak et al. (2024); Rohbeck et al. (2025), and previous unbalanced FM variants Eyring et (2024); Wang al. (2025). Our contributions can be summarized as follows:
et al. (2025a): Huguet et al. (2022); [Sun et al. (2025). balanced-distribuute
We propose WFR-FM, a novel framework that extends flow matching unbalanced distri- butions by jointly regressing transport velocity fields and growth rates. WFR-FM is entirely simulation-free, eliminating costly ODE integration and providing a scalable, robust, and OT-principled method for modeling unbalanced dynamics of population snapshots.
We establish theoretical guarantees for WFR-FM: minimizing its loss exactly recovers dy- namic unbalanced optimal transport under the WFR metric, i.e. the constructed cell popu- lation trajectories follow WFR geodesics.
We apply WFR-FM to handle multiple time points, making it particularly suitable for single-cell transcriptomics. Empirically, this yields more efficient and robust trajectory inference compared to existing neural differential equation and OT-based baselines.
2 RELATED WORKS
Unbalanced Optimal Transport and WFR Distance, Unbalanced optimal transport extends clas- sical OT Kantorovich (1942); Benamou & Brenier (2000) by either penalizing deviations from marginals (Benamou| 2003] Figalli 2010 Caffarelli & McCann 2010) or in its dynamic form by adding a penalty on source terms. A widely used variant is the Wasserstein Besserstein-Fisher-Bee Fisher-Rao (WFR) distance [iero al. (0016) Chizat al. (2018a): Kondratyev et al. (2016), which penalizes the
2

---

Published as a conference paper at ICLR 2026
squared growth rates and admits closed-fore solutions in special cases Chizat al. (2018a), with dynamic and static formulations shown to be equivalent (Chizat et al., 2018b Liero et al. 2018
Entropic regularization enables efficient solvers for static problem (Liero al. 2018 Chapel al.
2021]. Clancy & Suarez (2022) studied interpolating splines on the WFR space. Recently, deep learning approaches dynamic WFR have been applied to single-cell dynamics (Sha et al.) 2024 Peng et al. 2024| Zhang et al. 2025aa Sun et al.) (2025), but they depend on costly ODE simulation.
Action Matching (Neklyudov et al., 2023) and WLF (Neklyudov al., 2024) propose the non-flow- matching simulation- free method for WFR problem. Here WFR-FM introduces the simulation-free Now-matching-basse WFR solver that eliminates expensive trajectory integration in training.
Flow Matching and Unbalanced Extensions. Flow matching (Lipman et al. 2023 Liu et al.
2022
Pooladian et al. 2023 Albergo & Vanden-Eijnden] 2022) is a simulation-iree framework for
learning velocity fields via conditional probability paths in generative models, avoiding explicit tra-
jectory simulation. Subsequent work has expanded with OT coupling (Tong et al. 2024a Klein
2024), conditional generation (Rohbeck et al. 2025). muuli--arginnl formulations (Rohbeck ct al.
cet al.
2025 Lee et al. 2025), stochastic dynamics (Tong et al. 2024b Lee et al. (025), and other
generalizations (Kapusniak et al.) 2024] Zhang et al.) 2024b; Atanackovic al., [2025] Petrović
jet al. 2025). yet most assume normalized distributions, thereby conserving mass. Motivated by
real-world application such as cell differentiation with proliferation and apoptosis. unbalanced flow
matching has recently emerged (Byring et al.] [0244] [Ca et al. 2025 Corso et al. [0025] Wang et al.
20251. but existing approaches either regress only velocity fields (Eyring et al. [0244] Corso et al.
2025 Cao et al. 2025) or, when jointly modeling velocity and growth (Wang et al. 2025). still
rely on ODE simulation for post-training and the coupling mechanisms are not consistent with un-
balanced optimal transport. WFR-FM FM overcomes these limitations by providing an OT-principled formulation of unbalanced flow matching, which recovers the geodesics of the WFR metric.
Single-cell Trajectory Inference. Many algorithms have been proposed for single-cell trajectory inference in the multiple-time-ponnts setup, including static optimal transport (caseScchee al.) (201)) Klein et al. ); Halmos et (202)) |Banerjee et (2024). continuous formulations via neural ODEs Tong ct al. (2020); (H000)eHuuuuet et al. (2022): (Zhagg et al. (2024a): [Sha et al. (2024); Tong
ct al. (203)); Gu et al. (2025); Choi & Choi (2024), and stochastic variants Neklyudov et al.
20023
2024 Albergo et al. (2023); Maddu et al. 2024); Yeo et al. (202T); Chizat et al. (2022): Ventre
et al.
(2023): Lavenant et al. (2024): Shi et al. (2024) Koshizuka & Sato (2023): Bunne et al.
(023)) )
Chen et al. (022b): Diang& & Wan (024); Zhang et al. 2005a);: [Sun et al. (2025). More recently,
flow matching has been applied in this domain improve scalability and training stability (Tong
ct al. 2024a Rohbeck et al. 2025) Klein et al. 2024) Tong et al. 2024b Lee et al. [025 Kapusniak
al.
2024| Atanackovic et al., (025). The biological importance of proliferation and apoptosis has
further motivated extensions to unbalanced settings (Neklyudov et al. 2023 2024; Eyring et al.
Wang et al.] 2025). Notably, VGFM (Wang et al.] 2025) explicitly incorporating growth estimation into the flow-matchin framework. However, existing unbalanced approaches either rely on static approximations or require costly ODE simulations, or they lack a truly simulation-free flow matching framework that can jointly recover velocity and growth in single-cell dynamics.
2024
MATHEMATICAL BACKGROUND
Setup. Consider the distribution density function of two measures (a((2) and (1x) (also denoted No, 11) at times t = 0 and = 1 respectively defined over X' C Rd. Let M+) represent the set of all absolutely continuous finite measures support on X. The total mass of μi and μ may be different. We aim to learn a dynamical WFR-OT path connecting Mo and μ4 (defined below) via simulation- free framework. In this section and Section 4 we focus on the case of two measures. In Section 5
we will turn to the dynamic WFR problem with observation at multiple time points.
Dynamic Optimal Transport via Flow Matching, Dynamic OT provides a principled framework for interpolating between probability distributions by minimizing the kinetic energy of transport paths (Benamou& & Brenier 2000). In this setting, the transport is represented by a time-dependent probability path {Pe (Pe and a velocity field u, which jointly satisfy the continuity equation O,pt + V (p.u.) = 0. This dynamic formulation yields both natural interpolation between distribu- tions and a notion of geodesic distance. Recent work (Tong et al.) 2024a) connects this formulation to FM, which learns parameterized velocity field vo az) ) approximate the OT flow. The training
3

---

Published as a conference paper at ICLR 2026
objective is simple regression against the target field,
(F) Lr (l at) uz (z))))
avoiding costly backpropagation through ODE solvers. A key ingredient is the use of conditional probability paths. Instead of directly constructing the marginal path Pt, one defines conditionals pe(-Iz) with tractable vector fields ux(-1z). For example, given endpoint samples (æo, æ1). Gaus- sian bridge pr (a (o. x1) yields a closed-form vector field, and the marginal regression loss becomes
(CFM(() = B. (11z) (t x) [1]]
When conditioning pairs (Eg, 21) drawn from OT coupling (xv, x1). ) the resulting objective recovers the dynamic OT flow. Thus, the conditional path, conditional loss, and coupling (i.e. conditions) designs jointly determine if learned dynamics match OT geodesics general bridges.
Continuous measure flows. Consider a time-dependent measure path p: [0, 1] x Rdd - R__ a time dependent vector field [0, x IR - Rd and a time dependent rate function g [0, x| F R. Assuming enough regularity, this vector field and rate function defines unique time dependent flow Φ : [0, x Rd Rd, m [0, 1] x Fd - IR defined via an ODE:
dt (m) = (u(()(())))
d ( (c(x)
Φe(x)) x, mo( given.
(3.1)
(1b)(a)))))
which defines a 'push-forward" mapping with changing mass (Ø., m))+ MM (RE - (R
for each f. E [0, 1] such that
m₂(x)
Pr (x) (Øe, m))= po(Ø; (z))
(o)) det 'x);'(x)
(x)
mo( (x))
and satisfies the continuity equation with source term
(.p:(x) + V= (z)))))((z)) =: (x))
(3.2)
WFR Optimal transport. The dynamical unbalanced OT described by WFR metric (Chizat al. 2018a| b Liero et al. 2018) is defined as
inf ('' / [ [11lu(x:
p.g. Jo 2 WFR? 3 (Ho,
(3.3)
s.t.
0.p + V (pus = P9. Pa = Po: P1 11,
where 5 is hyperparameter in the WFR metric that balances the contributions of transport and birth-deate (or mass creation/annihillttoon)) terms in the action. To design the conditional path of flow matching for WFR, it is essential to study the geodesics between two mass points under such a metric. Indeed, the closed-form WFR distance between two Dirac measures mgd2, and m n2 is given by Chizat et al. (2018a) Liero et al. (2018)
I x1 12
WFR-DD (mo m1 28 28 m2 2 ht/m//mmmooooooooo cos
(3.4)
28
where cos( x)) cos(min })). When | - æılle < σō. the corresponding geodesic curve i.e. optimal transport path, named traveling Dirac is described by
(t) = A 2Bt + mo,
u(t)m(t) =
(3.5)
where A, B, wo can be computed as follows
mom
mam1
A mo m1 2
B mg
1 1+
V
(3.6)
mgm1
|||1 xolle
28T We
tan
1+
28
where I is the unit vector pointing from 2p #
4

---

Published as a conference paper at ICLR 2026
To construct the coupling in flow matching training, also refer to the equivalent static Kan- torovich form of WFR problem, which is also studied by Chizat et al. (2018b); Liero et al. (2018)
inf W vi) = (no ( (23))² / WFR-DD3 (no((, y)oz, Yr (a )ydyd,
(3.7)
where (7o 71) is called semi-coupling, and it satisfies the constraints
T(Ho, p1) def. {(ro.m)) E (M+(23)) J yo(a,, v)dy = po(a:), / n (2c, y)da =
(1 m1(1))
For fixed pair (ar, y), 7o0 (ar, y) represents the mass at the initial time sent from æ, while 71 (1 y) represents the corresponding mass at the final time received by y. Since the transport is unbalanced, the mass sent from æ may not be equal to the mass received by y. To efficiently solve such WFR coupling, is proved that the WFR problem is also equivalent to optimal problem (Chizat et al. 2018b
(OET)
Liero et al. 2018
1/e / -2 In cos
x yllk
(H, u1)
28²
inf
)r( 2)drdy
nEM, (22)
28
(3.8)
+ +KL( / 1(x y)dyllpo(xz)) + KL( / 7(x y dz
The relation between (70, 71) and y has been studied previously under general condition (Liero eet al. 2018). In this work, we also give straightforward proof (see Appendix (A.I) for the WFR problem that the optimal semi-coupling can be constructed from the OET coupling 7 as following.
Theorem 3.1. Let yy be the optimal coupling of the OET problem (.8). then the semi-coupling 7(,y)) (o y) Jx (l 7n (1x y) Jxx (==.)))= /n (y) solves the static WFR problem (37))
7(a.3)
EMMULATIO--RRES TRAINING FOR WFR PROBLEM
Given a vector field u, and rate function Sr that generate the measure path Pt by (3.2), we aim to learn time-dependent vector field va t) and time-dependent growth rate function go (æ. t) both parametrized by neural networks, with loss function specified by minimizing the intractable unbalanced flow matching objective
LUFM Φ) - -//-// (11oo(z,)) - (x(x) + (1).( (x ( (x)dd
(4.1)
Next, follow the idea of conditional flow matching (CFM) Lipman et (023) to learn ux.))) and ge (x) from conditional path, conditional loss and coupling guided by the WFR geometry.
4.1 CONDITIONAL PATH CONSTRUCTION
Conditional measure path Following the approach of defining conditional probability path in analogy to flow matching (Lipman et al. 2023). we define conditional measure path w.r.t the condition variable such that Pr (
m2(2), (z)), where the time-dependent conditional measure Pr (x)2) is decoupled into time dependent mass m,(2) and time-dependent conditional probability densityç: (x)))))) z). The un (x(2z) and (z) z) are conditional vector field and conditional rate function generating (=](2) from po(x) 2), i.e.
(xlz))
= pe (a: (z)2))))))
(4.2)
Marginal measure path Assuming z ~ q(2). we define the marginal measure path from the con- ditional measure path
Pe (x = |x (((a))))(()ddz
(4.3)
as well as the marginal vector field and rate function
(x)) (x) 2) (((m)))g(=) dz,
ge( = (x(2)))) (((l))((()))) dz pale)
(4.4)
(x)
5

---

Published as a conference paper at ICLR 2026
The following theorem sketches the relation between marginals uz(x), g((x) and conditionals (z(2( and (c(2))2) with the proof present in Appendix A.2
Theorem 4.1. The marginal vector field and rate function 4..4) generates the marginal measure path (433) from po(x) for any q(z) independent of a and t.
Conditional Gaussian measure path As convenient choice, we introduce the conditional Gaus- sian measure path (CGMP)
(2] = N NN m (2), (z)), Pr(zz)) - m, (2)
(4.5)
The conditional vector field and the rate function can be easily computed follows.
Proposition 4.1. For CGMP defined above 4.5 the conditional vector field and growth rate func- tion are
o((2) (z
(a m((=)) + (z), 9e z) = In (z
(4.6)
The proof is left to Appendix A.4 Note that the conditional vector field shares the same form as the conditional Gaussian path used in balanced FM (Lipman al. 2023). This allows algorithm to naturally include flow matching as special case by setting m, 2) 1.
4.2 CONDITIONAL Loss DESIGN
With designed (=/z)) and (x]lz), we propose regress them by minimizing the conditional unbalanced flow matching (CUFM) objective
(CUPM Φ) (==-0((011)===00)==-----
(x(22))1 + 19 (x(x))))))))))))))))((()))))(( (2)))) (4.7)
The crucial difference between the balanced CFM objective (Lipman et al. 2023) and CUFM ob- jective (4.7) is the time dependent mass term (2))). In the traditional FM, we only simulate the position of particles via the learned flow, while in unbalanced FM, the mass of the particle varies depending on time. The idea here is to weight the regression error by mass, which is assumed to be uniform in balanced setups. When (<r)))
the UFM objective 4 and CUFM (4.7) objective reduce to the well-known FM objective and CFM objective (Lipman et al. 2023).
The following theorem claims that the gradient of UFM and CUFM objective w.r. parameters (0, Φ) are equal, justifying the designed conditional loss function, with the proof present in Appendix A..3
Theorem 4.2. (x()) 0 for all X and [0, 1], and q(z) is independent of and t, then LUFM(0 (0, ②) = LCUFM(O, Φ) C, where C is independent to (0, Φ). Thus they have identical gradients w.r. (0, Φ), i.e.
VeLLUFFM(O.C Φ)= Ve.LCcUFM(0.0)) Φ).
4.3 COUPLING IN WFR-FM SOLVES THE DYNAMICAL WFR PROBLEM
Now, we consider two measure (x( and μn (æ) as source and target of a WFR problem, and the semi-coupling of the static form denoted as (70- 71). Without loss of generality, let fo( be probability density on X, then To (Co. æ1) is probability density on X²2 The conditional variable
(æo, 211 ~ Yu (Zg. 21), represents the position of the source and the target points. Since the
solution of dynamical WFR of two Dirac is traveling Dirac Chizat et al. (2018a), we choose special
CGMP named traveling Gaussian
ds
p=( xo. x1) = (2/x0 + wo
021 == -N(x/00 + wo x1), (0A)(20.21))
Jo m, (xo, x1))
(4.8)
=\(x><0 +
We
(arctan(
At B
arctan( VmoA- - B2 02))
-B
VinA= - B
VmoA- - B2
where wo and my (20, x1) are defined as 3.5p and 3.6 with boundary condition mo mo(
1, my (2o, 21) 71 (m).
=0( We set mo = 1 because m, describes the mass evolution relative to show that the marginal boundary
the initial mass along each trajectory. Under this setting, measures approach An and μn respectively as o - 0, and the induced measure path follows the WFR geodesic (proof in Appendix

---

Published as a conference paper at ICLR 2026
(a) Dynamic unbalanced optimal transport
(c) Kantorovich unbalanced optimal transport Eq. (3.7)
Eq. (3.3)
J mol v)dy = po(x)
(zt, 11dz (x
440
Pt. Ut 9t
Equiv alent
P1
Special Case
Condition Coupling
(b) Analytical Path from Dirac to Dirac
(d) Limulation-Fress Solution Process Condition Loss Eq. (4.7)
Eq. (3.4)- (3.6)
Condition Patn
Analytical solution?
Wt
(at, mp)
(Ze. mc)
(11,,
velocity
Ug (t, x))
growth
Pa (f,
Figure I: Overview of WFR-FM. (a) Dynamic unbalanced OT aims to interpolate intermediate state in way that minimizes the action. (b) Kantorovich unbalanced OT seeks the semi-coupling that minimizes total cost. (e) Analytical trajectory from a Dirac distribution to another Dirac distri- bution. (d) Jointly train the velocity field and the growth rate in the simulation- ffree manner.
Proposition 4.2. Take travelling Gaussian 14.8 CGMP, the marginal boundary measure Po(xo) + (z)). P1 (1(1)) + p1 (x1) as a + O. Also, the flow generated by travelling Gaus- sian recover the optimal path of dynamical WFR problem
3.3). thus the marginal ux(x) and gx(x)
solve the dynamical WFR problem between and μ41-
mi VI++z )f(1 - t) + m remains non-constant. Mass conservation only recovered as the growth penalty - oc, leading to the balanced OT limit characterized below (proof in Appendix A.6
Notably, even for balanced masses
the WFR mass trajectory m(t) = 2m
Proposition 4.3. With balanced source μi and target μı, let oo (3.3). then the travelling Dirac 3.5 reduces to uniform rectilinear motion with conserved mass, and the travelling Gaussian 4.8
reduces un (x|22 I Ep. 90 90(2) (z)=z) O, thus we recover OT-CFM (Tong et al. 2024a).
WFR-FM WORKFLOW FOR TRAJECTORY INFERENCE
In this section, we discussed the setup and workflow of WFR-FM in actual calculations (Figure such as reconstructing dynamics from ScRNA-seq data with multiple time points.
Multi-time point WFR Formulation. Given samples from distributions μt; at K + discrete time points, to, .i...... tk. our goal is to reconstruct the continuous evolution of these distributions. To this end, solve the WFR problem, which entails minimizing the WFR action functional subject to the constraint that p(x; t) coincides with the observed distributions ft, at the specified time points.
tK to inf J J& (Iuu(or 1)
WFR
- 19 0)1 px(x)) dæd/
Mtx
2
P.9.u
(5.1)
s.t.
Oxp + Va (pu)) = pg: Pto Moo -.o....... Ptx 14x
The multi-time problem 15.1 reduces to sequence of WFR steps, as established by the following proposition (proof in Appendix A.7). In practice, share vector fields and growth rates across all time intervals to enhance generalization.
Proposition 5.1. The solution the multi-time WFR problem 5.1 is equivalent to the concatena- tion of the solutions of successive time points.
WFR-FM based on Mini-batch WFR-OET. When solving the coupling, one needs to pre-compute the WFR-OET problem defined in (3.)) However, computing the full mapping directly on large- scale datasets expensive. Following OT-CFM (Tong al. 2024a), adopt mini-batch strategy.

---

Published as a conference paper at ICLR 2026
Ex dw; denote the empirical distribution supported the observed data points. We par- tition M4 and μr into the same number of respectively. For each pair (mD Mi ) we solve the WFR-OET problem to obtain an optimal cou- pling 700, The global coupling is constructed by concatenating these mini-batch solutions, i.e., Db-1 700), which approximates to the full coupling between Ao and M1- We include further discussions on the intuition, limitations, and practical considerations of using mini-batch WFR-OET in Appendix and provide empirical sensitivity analysis in Table
Let μ :E===E, =,
B, denoted by { ) and {pr (b) Ib=1
Complete Workflow We described the training workflow WFR-FM for trajectory inference problems with multiple time points through Algorithm I and inference workflow in Appendix E
Algorithm 1 WFR-FM Training Workflow
Require: Sample-able distributions Pto: Me,..... Me bandwidth o, training batch size b, batch size of the OET problem B, WFR penalty coefficient 6, vector field vy (v, L). growth rate go (x, t).
1: for k 0 K do
2:
((k) f- OET(H M or mini-batch (ET Moets B) (Defined in 3.8 7 (..)))
3:
(k) To (ax, y) & Ix y(0) (2, Me, (a), (k)
(a.
(x. (x. v) < Ix 2 (z.))d M (1) (Theorem 3.1
4: end for
5:
while Training do
6:
for == K do (k) (I1, Dle41) ~ 70 1e Sample batches of size i.i.d. from the datasets */
7:
8:
Calculate the constants A, B, we and that path depends on. (Defined in 3.6
9:
I L~ 1(0, 1),((k) e(K) + 14 (la+1 lt)
10:
Mea) + Din + woAr( )(+11) (Defined in 3.5 and 4.8
11:
(8) ~ N(nw, OFI) (Defined in 3.5 and 4.8
12:
(k) + ww/(me/ (Z Ltet. [fx+1 tx)) (Defined in 3.5 and 3.1
13:
(()) + at 12 Bie+i ) [tx+1 tx) (Defined in 3.5 and 3.1
14:
(*) + me( Dlas1 1/10
(Iez Dles1) (Defined in 3.5 and
4.8
end for
16:
Concatenate (0) (100, (t) g(())))))))))) m [
[==1 into batch tensors {xr 1, 1 y m²)
17:
CC Φ)= + (0 (9) u0 + (g t9) -g-(()) mm (Theorem 4.2
18:
O, Φ +- Update((B (f, Φ), (VoLcUrM 0, Φ), V LCUCF
end while
20:
return Vg and g
EXPERIMENT RESULTS
We evaluate WFR-FM with key questions Q1: Can WFR-FM transport Nio μt," ? Q2: Does it approximate the dynamic WFR solution? Q3: How accurate is interpolation of unobserved data? Q4: How does its scalability compare to existing methods? Q5: underlying birth-death dynamics?
How well does it recover the
Accuracy of Distribution and Mass Transportation (Q1) We evaluate the distribution time- propagation task on three synthetic datasets: Simulation Gene (Gene), Dygen and the 1000D Gaussian Mixtures (Gaussian) using two metrics (Appendix B): the I-Wasserstein distance (W1) to measure the similarity between normalized predicted and true distributions, and Relative Mass Error (RME) to assess how well the model captures cell population growth. WFR-FM consistently achieves superior performance across all datasets, with the lowest Wi and near-zero RME (Table I). In particular, it improves distributional matching on low- and mid-dimensional datasets and re- mains the most accurate in the dinensnnnnio case. We also evaluate WFR-FM on real datasets in Appendix B Visualization on the Gene dataset and the real-world embryoid bodies (EB) dataset MMonn et al. 2019) further confirms that WFR-FM captures accurate dynamics (Figure 2).
Closeness with Action of Static WER (Q2) We next examine whether WFR-FM provides a valid solution to the dynamic WFR formulation by comparing the path action of inferred trajectories with
8

---

Published as a conference paper at ICLR 2026
Table 1: Mean W1 and RME (only for unbalanced methods synthetic datasets. For the baselines that exhibit randomness in inference, we report the mean value and standard deviation over 5 runs.
Method
Gene (2D)
Dyngen (5D)
Gaussian (1000D)
Wi (L)
RME (U)
W1 (L)
RME (L)
W1 (4)
RME (J)
MMFM Rohbeck ct al. 2025
0,298
1.371
2.833
Metric FMI KKpusmakk el al. 2024
0.311
1.767
3,794
SF2M Tong cl al. 2024b
0.224
1.277
3.543_nmm
MIOFlow (Huguct ei al. 2022
0,148
0.965
2.858
TIGON (Sha al. 2024
0.045
0.014
0.512
0.047
2.263
0.127
DecpRUOT ZZhang al. 2025a
0.043
0.017
0.623
0.065
3.785
0.303
Var-RUOT' (Sun
2025
0.079
0.008=ax02
0.522=a.00
0.177
2.813=000
0.041
UOT-FM (Eyning
2026
0.093
0.010
1.204
0.097
2.771
0.033
VGFM (Wang al. 2025
0.046
0.006
0.598
0.037
3.010
0.037
WFR-FM (Our))
0.019
0.001
0.135
0.005
2.233
0.044
Gr
Gouun
74.0
Xo
Figure 2: Learned trajectories on the (a) Simulation Gene dataset (b) EB dataset
the reference action obtained from the static WFR-OET solver (Table 2). Across all datasets, WFR- FM consistently yields the closest approximation to the true action. While some methods such as Var-RUOT produce actions that are lower than the reference action and also below those of WFR- FM, they indeed trade off distribution fidelity to achieve a lower action. These results show that WFR-FM faithfully approximates the WFR geodesics in distribution interpolation task.
Interpolation Results under the Hold-One-Oue Out Experiment (Q3) We evaluate how well WFR- FM interpolates unobserved data in real multi-tima scRNA-seg datasets by holding out one inter- mediate time point during training (Table 3). Across EMT (Cook Vanderhyden, 2020), CITE- seq (Lance al. 2022), and mouse hematopoiesis data (Mouse) (Weinreb al. (0220), WFR-FM achieves the best interpolation accuracy, while comparable with top baselines on EB dataset. This suggests that unbalanced methods may better accommodate these datasets, possibly due popula- tion imbalance from cell proliferation; see Appendix B. for further discussion.
Scalability Large-Scale Dataset (Q4) We examine the scalability of WFR-FM on the large- scale EB dataset (100D) by comparing runtime, memory, and accuracy (Figure 3 WFR-FM achieves both high accuracy and efficiency, outperforming methods that only learn velocity fields
Table Path action on synthetic datasets.
Method
Gene
Dyngen
Gaussian
TIGON
1,367
11.654
6.794
DeepRUOT
1.286+
14.070+aast
14.370+1106
Var-RUOT
1.208 =0m22
8.357
7.146=088%
WFR-FM
1.305
9.410
7.529
Static Reference
1.333
9.569
10.531

---

Published as a conference paper at ICLR 2026
Table Mean W1 over held-out time points on EMT, EB, CITE and Mouse datasets.
Method
EMT (10D)
EB (50D)
CITE (50D)
Mouse (50D)
MMFM
0.323
11.213
38.521
8.263
Metric FM
0.314
10.726
37.342
7.753
SF2M
0.308
10.986=0.06
38.333
8. 8.646=0004
MIOFlow 0.325
10.960
39.574
7.779
TIGON
0,360
11.080
38.159
6.868
DeepRUOT
0.323
10.075
37.892
6.847
Yar-RUOT
0.320
11,035 000017
38.393 =0:29
8.672+0000
UOT-FM
0.322
11.344
38.649
9.332
VGEM
0.301
10.370
37.386
8.496
WFR-FM 0.298
10.157
37.221
6.586
in accuracy and ODE-based methods in efficiency. WFR-FM provides favorable balance be- tween computational cost and accuracy, making well-suited for large-scale single-cell datasets. Sensitivity analysi confirms WFR-FM's robust- ness to key hyperparameters, specifically the growth penalty d (Table 5) and batch size for mini-batch WFR-OET (Table 8).
MMFM
13.5
Ny Distance
UMT-FM
T6ON
Memory (GB)
(log scall)
Evaluation of Learned Growth Dynamics (Q5) To confirm WFR-FM captures birth-death dynamics beyond marginal matching, we eval- uate the learned growth rate (x, f) on the Simulation Gene dataset, where ground truth is provided by known SDEs. We calculate the Pearson correlation coefficient (fcor) between the predicted growth rate and the ground truth X²
Var-UUm
12.5
WOP
12.0
MFR VGFM
DeopRUCT
11.0
0.1
200
400
600
800
1000
Training Time (s)
Figure 3: Efficiency on the 100D EB dataset
growth rate = Og 지 on the data points from the dataset. Table shows that WFR-FM outper- forms baselines, validating its ability to reliably estimate growth rates.
Table 4: Growth rate correlation (geon) on Simulation Gene dataset. Best results are in bold.
Method
TIGON
DeepRUOT
Var-RUOT
Action Matching
WFR-FM (Ours)
Correlation (gcorr)
0.9705
0.9688
0.9214
0.5851
0.9913
CONCLUSION, LIMITATION AND DISCUSSION
In this work, we introduce WFR-FM, a framework unifying the WFR metric with flow matching for simulation-free, dynamic unbalanced OT. By jointly learning learning velocity fields and growth rates, WFR-FM provides a principled formulation with theoretical guarantees that exactly solves the WFR problem. Extensive experiments demonstrate its effectiveness in reconstructing dynamics with proliferation and apoptosis, establishing it as a robust tool for biological trajectory inference.
Like other OT-FM methods, WFR-FM relies on solving static OT problem, which can be costly for very large datasets. To address this, we adopt approximation strategies such as mini-batch OT Sinkhorn iterations. Incorporating uncertainty quantification could also enhance the performance on noisy single-cell datasets. While our focus has been on WFR geometry, the framework established is indeed general: as long as the static OT problem can be efficiently solved and the trrcce path admits a closed-form solution (e.g., when linear growth penalties are chosen, as Chizat et al. (0188)), our approach can be extended to other unbalanced transport functionals. We anticipate that this generality will enable classes of mawchat methods, broadening their applications across scientific domains where mass is not conserved other machine learning application scenarios with unbalanced datasets.
and entropy-regularive
WFR-FM
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
We have taken several measures ensure the reproducibility of our work. All theoretical results accompanied by complete proof' provided in the Appendix. Detailed descriptions of data processing steps are also included the Appendix to facilitate replication of our experiments. The source code is available at https:: /tgtttb....t//gg////.///.g.... com/ htt///aaanwwleeaaeee i Peng/ WFR-FM
AUTHOR CONTRIBUTIONS
Q.P. P.Z., and T.L.conceived the project. Q.P. and Z.W. implemented the algorithm and conducted data analysis. J.Y. and Y.S. performed the mathematical framework design and theoretical analysis. Q.P., Z.W., J.Y., and Y.S. contributed equally to this work. P.Z., T.L.,, L.Z., and Q.N. supervised the research. P.Q., Z., J.Y., Y.S., and P.Z. wrote the manuscript. All authors revised and approved the manuscript.
ACKNOWLEDGMENTS
This work was supported by the National Key R&D Program of China (No. 2021 YFA 100301330 T.L., and 2024YFA091 9500 L.Z.) and National Natural Science Foundation of China (NSFC No. 121 to T.L. & P.Z. & L.Z., 12225102 to L.Z., 8206 100646 to P.Z., and T2321001 to P.Z. & I.Z.).. We acknowledge the support from the High-performance Computing Platform of Peking University for computation. We thank the anonymous referees for their valuable feedback and constructive suggestions.
REFERENCES
Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic inter- polants. arXiv preprint arXhv:2209.1557., 2022.
Michael Albergo, Nicholas M Boffi, and Eric Vanden- Stochastic interpolants: uniffynng framework for flows and diffusions. arXiv preprint (tt.::22330077997, 2023.
Lazar Atanackovic, Xi Zhang, Brandon Amos, Mathieu Blanchette, Leo J Lee, Yoshua Bengio, Alexander Tong, and Kirill Neklyudov. Meta flow matching: Integrating vector fields on the wasserstein manifold. In The Thirteenth International Conference on Learning Representations, 2025.
Amartya Banerjee, Harlin Lee, Nir Sharon, and Caroline Moosmåšller. Efficient trajectory infer- ence wasserstein space using consecutive averaging. arXiv preprint arXiv:2405. 19679, 2024.
Jean-David Benamou. Numerical resolution an "uunbalanced" mass transport problem. ESAIM: Mathematical Modelling and Numerical Analysis, 37(5):851-866 2003.
2003.
Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution the monge- kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-393, 2000.
Charlotte Bunne, Ya-Ping Hsieh, Marco Cuturi, and Andreas Krause. The schrödinger bridge be- tween gaussian measures has closed form. In International Conference on Artificial Intelligence and Statistics, pp.
PMLR, 2023.
Charlotte Bunne, Geoffrey Schiebinger, Andreas Krause, Aviv Regev, and Marco Cuturi. Optimal transport for single-cell and spatial omics. Nature Reviews Methods Primers, 4(1):58, 2024.
Luis A Caffarelli and Robert McCann. Free boundaries in optimal transport and monge-ampere obstacle problems. Annals of mathematics, pp. 673-730, 2010.
Robrecht Cannoodt, Wouter Saelens, Louise Deconinck, and Yvan Saeys. Spearheading future omics analyses using dyngen, multi-modal simulator of single cells. Nature communications, 12(1))3942, 2021.
Zihan Cao, Yu Zhong, and Liang-Jian Deng. Taming flow matching with unbalanced optimal trans- port into fast pansharpening. arXiv preprint 1XX.:2503..4975 2025.

---

Published as a conference paper at ICLR 2026
Lactitia Chapel. Rémi Flamary, Haoran Wu, Cédric Févotte, and Gilles Gasso. Unbalanced opti- mal transport through non negaatiev penalized linear regression. Advances in Neural Information Processing Systems, 14:23270-23288 2021.
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 nanoball arrays. Cell, 185(10):1777-1792, 2022a.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou. Likelihood training of schrödinger bridge using forward-backwan SDEs theory. In International Conference on Learning Represen- fations, 2022b.
Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavits Vialard. An interpolating distance between optimal transport and fisher-rao metrics. Foundations of Computational Math- ematics, 18(1):1-44, 2018a.
Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard. Unbalanced opti- mal transport: Dynamic and kantorovich formulations. Journal of Functional Analysis, 274(11): 3090-3123, 2018b.
2018B..
Lénaïc Chizat, Stephen Zhang, Matthieu Heitz, and Geoffrey Schiebinger. Trajectory inference via mean-field langevin in path space, Advances in Neural Information Processing Systems, 35: 16731-16742, 2022.
Jaemoo Choi and Jaewoong Choi. Scalable simulation-free entropic unbalanced optimal transport. arXiv preprint arXiv:2410.02656 2024.
Julien Clancy and Felipe Suarez. 2022.
splines. arXiv preprint arXiv:2203. 15728,
David P Cook and Barbara C Vanderhyden. Context specificity of the emt transcriptional response.
Nature communications, 11(1):2142, 2020.
Gabriele Corso, Vignesh Ram Somnath, Noah Getz, Regina Barzilay, Tommi Jaakkola, and Andreas Krause. Composing unbalanced flows for flexible docking and relaxation. In The Thirteenth International Conference on Learning Representations, 2025.
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. In The Twelfth International Conference Learning Representations, 2024.
Alessio Figalli. The optimal partial transport problem. Archive for rational mechanics and analysis, 195(2):533-560, 2010.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z Alaya, Aurélie Boisbunon, Stanis- las Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, et al. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78): 2021.
Anming Gu, Edward Chien, and Kristjan Greenewald. Partially observed trajectory inference using optimal transport and a dynamics prior. In The Thirteenth International Conference on Learning Representations, 2025.
Peter Halmos, Xinhao Liu, Julian Gold, Feng Chen, Li Ding, and Benjamin Raphael. Dest-ot: Alignment of spatiotemporal transcriptomies data. Cell Systems, 16(2), 2025.
Guillaume Huguet, Daniel Sumner Magruder, Alexander Tong, Oluwadamilola Fasina, Manik Kuchroo, Guy Wolf, and Smita Krishnaswamy. Manifold interpolating trinsalt flows for trajectory inference. Advances in neural information processing systems, 35:29705-2977 2022.
12

---

Published as a conference paper at ICLR 2026
Qi Jiang and Lin Wan. A physics-informet neural SDE network for learning cellular dynamics from time-series scRNA-seq data. Bioinformatics, 1iiii20--i1227, 09 2024. ISSN 1367-4811.
Leonid Kantorovich. On the translocation of masses. In Dokl, Akad. Nauk. USSR (NS), volume 37, pp. 199-201, 1942.
Kacper Kapusniak, Peter Potaptchik, Teodora Reu, Leo Zhang, Alexander Tong, Michael Bronstein, Joey Bose, and Francesco Di Giovanni. Metric flow matching for smooth interpolations on the data manifold. Advances in Neural Information Processing Systems, 37:135011-135042, 2024.
Dominik Klein, Théo Uscidda, Fabian Theis, and Marco Cuturi. Genot: Entropic (gromov) wasser- stein flow matching with applications to single-cell genomics. Advances in Neural Information Processing Systems, 37: 17:107887-113944, 2024.
Dominik Klein, Giovanni Palla, Marius Lange, Michal Klein, Zoe Piran, Manuel Gander, Lactitia Michael Sterr, Lama Saber, Changying Jing, al. Mapping cells through time and space with moscot. Nature, pp. 1-11,0022
Stanislav Kondratyev, Léonard Monsaingeon, and Dmitry Vorotnikov. A new optimal transport distance on the space of finite radon measures. arXiv preprint (rXiv:1505.07740 2016.
Takeshi Koshizuka and Issei Sato. Neural lagrangian schrödinger bridge: Diffusion modeling for
population dynamics. In The Eleventh International Conference on Learning Representations, 2023.
Christopher Lance, Malte D Luecken, Daniel B Burkhardt, Robrecht Cannoodt, Pia Rautenstrauch, Anna Laddach, Aidyn Ubingazhibov, Zhi-Jie Cao, Kaiwen Deng, Sumeer Khan, et al. Multimodal single cell data integration challenge: results and lessons learned. BioRxiv, pp. 2022-04, 2022.
Hugo Lavenant, Stephen Zhang, Young-Heon Kim, Geoffrey Schiebinger, al. Toward mathe- matical theory of trajectory inference. The Annals of Applied Probability,
-0000 2024.
Justin Lee, Behnaz Moradijamei, and Heman Shakeri. high-dimensional snapshot data irregular time points. arXiv preprint arXiv:2508.0433 043511, 2025.
stochastic flow matching for
Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal transport in competition with re-
action: The hellinger-kantorovich distance and geodesic curves. SIAM Journal on Mathematical Analysis, 48(4):2869-291 2016.
Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal entropy-transper problems and new hellinger-kantorovich
kantorovich distance between positive measures. Inventiones mathematicae, 211 (3):969-111) 2018.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Repre- sentations, 2023.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0300 2022.
Suryanarayana Maddu, Victor Chardès, Michael Shelley, al. Inferring biological processes with intrinsic noise from cross-sectional data. arXiv preprint arXiv:2410.07500, 2024.
Kevin R Moon, David Van Dijk, Zheng Wang, Scott Gigante, Daniel Burkhardt, William S Chen, Kristina Yim, Antonia den Elzen, Matthew J Hir,, Ronald R Coifman, al. Visualizing structure and transitions in high-dimensional biological data. Nature biotechmology, 37(12))1482- 1492, 2019.
Kirill Neklyudov, Rob Brekelmans, Daniel Severo, and Alireza Makhzani. Action matching: Learn- ing stochastic dynamics from samples. In International conference on machine learning, pp- PMLR, 2023.
Kirill Neklyudov, Rob Brekelmans, Alexander Tong, Lazar Atanackovic, Qiang Liu, and Alireza Makbzani. A computational framework for solving wasserstein lagrangian flows. In Forty-first International Conference Machine Learning, 2024.
13

---

Published as a conference paper at ICLR 2026
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary Vetito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
Qiangwei Peng, Peijie Zhou, and Tiejun Li. stver: Spatiotemporal dynamics of single cells. bioRxiv, pp. 2024-06, 06, 2024.
Katarina Petrovic, Lazar Atanackovic, Kacper Kapusniak, Michael M. Bronstein, Joey Bose, and Alexander Tong. Curly flow matching for learning non-gradient field dynamies. In Learning Meaningful Representations of Life (LMRL) Workshop at ICLR 2025, 2025.
Aram-Alexandre Pooladian, Heli man, and Ricky TQ Chen. Multisample flow matching: Straightening flows with minibatch cou- plings. arXiv preprint arXiv:2304. 14772, 2023.
Carles Domingo-Enrion Brandon Amos, Yaron Lip-
Martin Rohbeck, Edward De Brouwer, Charlotte Bunne, 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.
Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal-transpon analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell. 176(4):928-94) 2019.
Yutong Sha, Yuchi Qiu, Peijie Zhou, and Qing Nie. Reconstructing growth and dynamic trajectories from single-cell transcriptomics data. Nature Machine Intelligence, 6(1):25-39, 2024.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching, Advances in Neural Information Processing Systems, 36, 2024.
Max Sommerfeld, Jöm Schrieber, Yoav Zemel, and Axel Munk. Optimal transport: Fast proba- bilistic approximation with exact solvers. Journal of Machine Learning Research, 20(105):1-23, 2019.
Yuhao Sun, Zhenyi Zhang, Zihan Wang, Tiejun Li, and Peijie Zhou. Variational regularized unbal- anced optimal transport: Single network, least action. arXiv preprint arXiv:2505.11823, 2025.
Panagiotis
Theodoropoulos, Augustinos D Saravanos, Evangelos Theodorou, and Guan- Horng Liu. Momentum multi-marginal schr\" odinger bridge matching. arXiv preprint arXiv:2506.10160 2025.
Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy, Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In International conference on machine learning. pp. 9526-9536. PMLR, 2020.
Alexander Tong, Manik Kuchroo, Shabarni Gupta, Aarthi Venkat, Beatriz P San Juan, Laura Rangel, Brandon Zhu, John G Lock, Christine I. Chaffer, and Smita Krishnaswamy. Learning transcrip- tional and regulatory dynamics driving cancer cell plasticity using neural ode-based optimal trans- port. bioRxiv, pp. 2023-03, 2023.
Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 2024a. ISSN 2835-8856. Expert Certification.
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, pp. 1279-1287. PMLR, 2024b.
Elias Ventre, Aden Forrow, Nitya Gadhiwala, Parijat Chakraborty, Omer Angel, and Geoffrey Schiebinger. Trajectory inference for branching sde model of cell differentiation. arXiv preprint AXxiv2307.006877 2023.
14

---

Published as a conference paper at ICLR 2026
Dongyi Wang, Yuanwei Jiang, Zhenyi Zhang, Xiang Gu, Peijie Zhou, and Jian Sun. Joint velocity- growth flow matching for single-cell dynamics modeling. arXiv preprint arXiv:2505. 13413, 2025.
Caleb Weinreb, Alejo Dodriguez-FFraticlli Fernando D Camargo, and Allon M Klein. Lineage tracing on transcriptional landscapes links state to fate during differentiation. Science, 367(6479): eaaw3381, 2020.
Hanshu Yan, Jiawei Du, Vincent YF Tan, and Jiashi Feng. On robustness of neural ordinary differ- ential equations. arXiv preprint arXiv:1910.05513, 2019.
Grace Hui Ting Yco, Sachit Saksena, and David K Gifford. Generative modeling of single- cell time series with prescient enables prediction of cell trajectories with interventions. Nature communications, 12(1):3222, 2021.
Jiaqi Zhang, Erica Larschan, Jeremy Bigness, and Ritambhara Singh. scNODE: generative model for temporal single cell transcriptomic data prediction. Bioinformatics, 406Supplement 146- i1154, 09 2024a. ISSN 1367-4811.
Xi Nicole Zhang, Yuan Pu, Yuki Kawamura, Andrew Loza, Yoshua Bengio, Dennis Shung, and Alexander Tong. Trajectory flow matching with applications clinical time series modelling. Advances in Neural Information Processing Systems, 198-0229840,7224, 2024b.
Zhenyi Zhang, Tiejun Li, and Peijie Zhou. Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport. In The Thirteenth International Conference on Lear- ing Representations, 2025a.
Zhenyi Zhang, Yuhao Sun, Qiangwei Peng. Tiejun Li, and Peijie Zhou. Integrating dynamical systems modeling with spatiotemporal serna-seq data analysis. Entropy, 27(5), 2025b. ISSN 1099-4300.
Grace XY Zheng, Jessica Terry, Phillip Belgrader, Paul Ryvkin, Zachary Bent, Ryan Wilson, Solongo B Ziraldo, Tobias D Wheeler, Geoff P McDermott, Junjie Zhu, al. Massively parallel digital transcriptional profiling of single cells. Nature communications, 8(1):14049, 2017.
15

---

Published as a conference paper at ICLR 2026
A PROOFS
A.1 PROOF OF THEOREM 3.1
Lemma A.1. The static form or WFR problem 3.7 is convex with respect to (a y) and yn (xx y).
Proof.
First, we prove that the integrand:
x
I (70, 71, x, y) = 72 + To cos
(A.1)
28
is convex with respect to to and 7/1 To do this, compute the Hessian matrix:
|
2 24
|| 26
$
2 cos
70
cos
H
(A.2)
la
-
l
COS
26
7
71
2
cos
26
70
71
Next, we check its principal minors. Since Yo. 71 of the matrix are non-negative. And for the determinant:
2
0
and cos
2 vl
0, the diagonal elements
28
det(H) cos? ||
y/l
| yl1
T0
71
cos2
To
(A.3)
28
28
Therefore, is a positive semidefinite matrix, and o is convex with respect to To and 7/- Our optimization objective:
I F/20,70 nl = fre I (voo2, y), (xx y), x, y) dadd
1
inf Tho (ll
25
70.
(A.4)
is a functional of "ro and 71. Consider any N0.a: 70,6 M1, n.b E (x and t []:
/ (tro, + t),,,, + (1 t),,)) daady
F a + (1 t0,0,,,,, nn + (1 11mmm
x²
< (1 (1o,ar (1,)) + (1-))
1) 1 (70.6: 71,6)) dady
x2
1 (70.0s + 1) (70.6, m1.1
(A.5)
Therefore, the problem is convex with respect to %o and m-
Theorem 3.1.
Let re" be the optimal coupling of the OET problem (38 then the semi-coupling (x..) Jo( 8) Tx (m Ho(x), (xx 4) (z
"'(xx))
Mi (y) solves the static form of WFR prob-
lem
3.7
Proof.
Let
I (ro(2, y), (x,
WFR g) 71 (),))
28
The computation of the WFR distance between two distributions can be formulated as the optimiza- tion problem (3.8):
inf (1o.)) / I (To v), 7 (x., y), x, y)dwdy
W (1)
(A.6)
25
20 = = / (ac, (2. pu m1(1))
s.t.
(A.7)
16

---

Published as a conference paper at ICLR 2026
By introducing the dual variables ((x), (b),, it can be transformed into the dual problem:
1 WHo,(HO:H4 (1) = sup
J I (rol(,, 7 y). æ. yddrdy
inf (oo
28
Jx d(x) 10(1 70(t y)dy -o(((0))))) do v(v) (Sx ri (a ydd 1 (1())) dy
/ (())) ((1))1)))d
sup (o.)))C(X))
x
211
s.t. ((x) S 1, (1y) S 11 - ((x))(1 - vb))) 2 cos
(A.8)
28
By change of variables, let u - log(1 Φ), logg1 -,),, the above dual problem becomes:
sup ( (1- (x=xp((uuz)))pp(t= + f(1 cxp(- (xx)
282 WFR
(A.9)
|
2112
s.t. u(x) + u S log
cos
(A.10)
26
The optimal solution (2, y), (x, y) of this optimization problem and the optimal dual variables u (x), U* (y) should satisfy the following KKT conditions:
(1) Complementary slackness condition:
(1 - exp((x"xx))) (/x 76 yddy mo(x)) 0
(A.11)
(1-exp(-o"(x)) (Jx "1 v)da - mi(t))
(2) stationary condition:
Oyo I (70, n, x. v)l exp(u (a)))
(A.12)
Oys I (7o, Mi, y)h (1 -1- exp(e'
=11-exp(o" (y))
(x y)
(x (x.v) y)
The solution we constructed, (a, y)
7 ""x 2 2)d2 (a(a)), 71 (ax y) Jx 7- y) 11 (1))),
satisfies the complementary slackness condition. Next, we need to verify that it satisfies the station- ary condition. Consider the equivalent unconstrained optimization problem (3.9):
1 WFR 11) 282 (x3+
(Jez -2 cos
||
yl12
inf
((x y ydddy
28
28
( 7(0; wddylllol(s) + KL
(/g 7(0, ()ddlllul())
KL
Transform into a constrained optimization problem:
1 WFR H1) 282
(fee -2 In
| -x0 v|1a
inf 7EM+(x²3
7(0 ydddy
28
28
(A.13)
+ KL (L KL (
s.t. / 7(0, v)dy =po(z), / 7(x, ve)ev=xd(pt
(A.14)
17

---

Published as a conference paper at ICLR 2026
Introduce dual variables u(æ)), v(y) obtain the dual problem:
1
inf (a.)) Po (M).EM (x2) (Jrr -2 In cos
(1 u1l 20
||
WFR ( u1)
sup
7(x, yddady
28
+ KL KL (L00))))) KL (p
- -fx u(x) (L 720, yddy - po( dx -
(ainn-((())) / o(y) (/x 7(2, y)dx - p1(1)) ((a
dy
x
inf (u,o)E0(ar (2 /
1 [1-yyl1 y1l2
- sup
- cos
uuæ)) - -v(y)
20
+ (KLpolluo)) + / ((x))o((z)ddr)) + (KL(pillun))
f u(x)) (m))d
(A.15)
Now, Po: P y are in three separate terms, so when solving the inner optimization problem, we can minimize these three terms separately. From the first term, obtain the constraint the dual variables:
s.t. ((x)) + (ty) I
log
cos
(A.16)
25
By solving the inner optimization problem, we can obtain the relationship between the dual variables and the primal variables:
| U1l2 )) u(xx v(y)
(xx 3) supp(?(<, (z. y)). log
cos
(A.17)
26
Po (x ( exp(-u(x)),
Pi (1 cxp(-o(u))
(A.18)
(a (x)
#1 (v)
Substituting these relationships into the dual problem of (3.9), we get:
[(1 exp(-u())))
f11 exp( (i-eep(-oou)))pn (())))
(F,
sup (u,)))C(())
28
(A.19)
This is exactly the form of the dual problem of (3.8) after the change of variables. Therefore, (3.8)
and (3.9) have the same dual problem, and the u(x), (, v(y) we introduced when constructing the dual problem of (3.9) are the same u(x), v(y) in the dual problem of (3.8). Consider the relationship between the primal and dual variables when both are optimal in (3.9):
x
))
V(x, y) E supp(?(xx, y)), - log COs
-x(x)) - 0- (3) =0
(A.20)
28
Jx yddy
Jx (x, y)dz
exp(-u'(3))
(A.21)
(x
M1(3)
Since:
(x,, y) (z, y)
|| 1/l2
a ha
cos
(A.22)
28
7" (x. y)
Substituting the solution we constructed, (0 y)
Jx (x, 2z, z)dz Ho( m (æ.
(xx y)
Ix (2, y)dz 11 (v), and the above relationship between primal and dual variables, we can ify:
Dy I hs = exp (0 n² (x))) exp 1 (1) + +o
=1 - (epp((u(xx)
Thus, the stationary condition is satisfied. The other stationary condition can be verified by sub- stitution in the same way. Therefore, the solution constructed satisfies the KKT conditions of problem 3.7 Accroding to Lemma A.I problem 3.7 is convex with respect to 7o- n. the solution
constructed is the optimal solution of 3.7
18

---

Published as a conference paper at ICLR 2026
A.2 PROOF OF THEOREM 4.1
Theorem 4.1, The marginal vector field and rate function 4.4 generates the marginal measure path 4.3 from po(æ) for any q(z) independent of a and t.
Proof. To verify this, it is sufficient to show that the marginal measure pe(a)) satisfies the continuity equation with source term B.2b
d 0sp((0)= dt / (=(x)=))(=)dz
(,p)(a))))(2)d)
Vx (ua_(((())))))(()))))))))) + (((00)))((())))()))
V / (((0))))((0))))()))) [((=0)))((0()))(()))
V. (p(( / (11z) (=(
dz) + p((x)) / 9 (2]2z) ( (((2))))(())
dz
px(
pr(x)
Vx (x)( (z))))
In (1), we use the continuity equation with source term of the conditional measure (=(a)))) (.2 In (2), we use the definition of marginal vector field and rate function (4.4). Thus, the marginal measure (xr) satisfies the continuity equation with source term (3.2). The proof is completed.
A.3 PROOF OF THEOREM 4.2
Theorem 4.2 If p(()) > 0 for all a X and [0,1], and q(z) is independent of x and t,
then LUPM (b, LC Φ) + C. where C is independent to (0, Φ)) Thus they have identical gradients w.E. (0, Φ), iee.
Vo.LLmul(0) = Vo. LCurM(()))
Proof. First, recall the two objective
(0 =Emum(o.1 / (Ilvo(x,))
(z)))) (2, t) (e(z))
LCUMM(0 () x==
(=(=))]1 (2, t) (1(
19

---

Published as a conference paper at ICLR 2026
By direct computation, have
O0 Φ) = =Bn0e000111 / (1lvo(2e = 2(v0(x (z)))
+ K1E1-e / (1190((z, t)1l (z t),
2 Beneetoun1 / (110 (e)
2 / (00(2, t), (atp///uuaaaaaalctapp(ddadaaaa
KIEm / (1loo(0,
2 (90( t), [((ak)))p(((333(()))))))}} -/oooo/el////ott/aaeee2330--0---
Beme(10n11 [xe] [1loo(a,,))] hrta:://-eakkoaaaauuaaa/aaaaaaad
+ xlBeneet0.11 / (1196(2 1)1 12gp
C
12m22(0,11 10oo::0..ll-mmeeeee---..::::: t) 1 2 t),
+ -0t- 2
()))
= LCUFM (0, Φ) +
where C' represents a constant independent to (0, Φb. In (1) we expand the marginal measure (=
by its difinition (4.3). (2), we decouple the measure Pe (2: (x) into mass m2(2) and density (r In (3), we rewrite the integral in the form of expectation, moving the two probability density ((2)
and fa( z) to the expectation operator.
Since the constant C is independent to (0, ()) we have
V Lu, = Vo, CCLFM
The proof is completed.
A.4 PROOF OF PROPOSITION 4.1
Proposition 4.1 For CGMP defined as 4.5). if m²(2) 0 for all t. then the conditional vector
field and rate function
01(2)
u: (2(2))
(x m1(2)) + n1(2z)
(z
(A.23)
(x(2))) = 0 In (z)
Proof. The CMGP satisfies the following continuity equation with source term (4.2
(xlz)) + V.
= Pt (=)=))
expand the conditional measure into mass and density, the equation become
m=(2)04p.<z 2) (2), + my Vz) (zx( m. (1))
(A.24)
Since m; m;(2) > for all f, we devide the both sides the equation by m, and get
0.7.( 2) + 0 In (z)p) Va (x)) Pr (
(A.25)
We further reorganize it to
(xl=)) Va (x< ( (a)l
a In mx))))
(A.26)
20

---

Published as a conference paper at ICLR 2026
Set 90 (x](2) 0 In (z),,, the equation of the density Pe (2r) (2) reduces to the continuity equation
( + Vx (u=(z)z)) [0)===))))))))))))))) 0
(A.27)
Note that this equation is exactly the continuity equation of the conditional Gaussian path in con- ditional flow matching. Thus, the conditional vector field of CGMP shares the same form with the conditional vector field of conditional Gaussian path in balanced conditional flow matching, which The proof completed.
(=) (,(==)) (z) x (=)) (2).
A.5 PROOF OF PROPOSITION 4.2
Proposition 4.2 Take traveling Gaussian (4.8 as CGMP, the marginal boundary measure
Po(xo) (o (o))) p1(x1) μ4 (x1)) + O. Also, the flow generated by traveling Gaus- sian recover the optimal path of dynamical WFR problem (33). thus the marginal ux(x) and ga
solve the dynamical WFR problem between μi and μ1-
Proof. First, we expand the marginal measure
Pc(x) = (c(x))))(2)dz=
Set / = 0, we get the marginal boundary measure
/ Po(z)2))(()ddz
Po(x:)
/ (n(z)z)moled 21)4 æn dx
N(x o' 1 70 (wo, æ: d
x2
(A.28)
{ M(x (o 21 d
X²
N(e wo, o 1)oo(2)))dn x
p4 * (N o2 = (x())
Set / = 1, we get the marginal boundary measure
/
p1(x)
pi (z))2))(2)dz
(1(012)) (Io, x (x00 æi dd
x²
/ N(xaa, 0²) 71 (2 21
7o (wo, a )dzoda:
70(20, x1)
(A.29)
Xx2 M(xx/11 2 I)m (Ia, 21 dæ
x
= * (0 (z) = ui (x)
Thus the learned unbalanced flow matches the source measure μu and the target source P1-
Next, show that the induced marginal path {pt), which is the push forward of traveling Dirac on "i (x, y). indeed satisfies the geodesic of WFR metric. The WFR distance between the marginal measure Ps and Pt (s < f) can be estimated follows.
21

---

Published as a conference paper at ICLR 2026
inf
/xa WFR (a, y) 7i" y)oy dxddy
WFR (m, (x, Mr yo. 71 71 (m. (m. m y, yo. 11)0 daady drdy
x22
(2)
t. s| WFR (70 y) (ax y) y drdd
x2
1t
8 @WFR
(A.30)
where zt (æ, 70. 71), mz( y. Yo. 71) are the position and mass at time t induced by the traveling Dirac connecting 7o y) and 71 (x. y). In (1) we replace the optimal semi-coupling between ps and Pt. namely (70" 71"), with a specific semi-coupling induced by the push forward of traveling Dirac on "ro (z, y). (2), we use the constant speed property of the traveling Dirac as geodesic under WFR metric, which has been proved by (Chizat et al. (018a) and (Liero et al. (0066). One can also check the constant speed property traveling Dirac by direct calculation. Utilizing the inequality above, we have
WFR (Ho: 11)
WFR Pa) + (P..P) + WFR (pe
1 WV
0) WFR (1) + WFR (1) + WFR p1)
WFRs( m1)
(A.31)
In (1) we use the triangle inequality of WFR, since it is metric on (+ (Chizat et al. 201 2018a).
In (2) we use (A.30) on time intervals [0, [s, and [t, IJ. In (3) we simply use the fact that
S 8 < L 1. Thus all the inequalities in (A31)) are actually equalities. Specifically, we have
WFR.(P..P.) = It - s/WFRs(Ao.H4)
(A.32)
This indicates that the marginal measure path {p:} is geodesic connecting the source μo and the target μ1 under the WFR metric. Thus, the corresponding marginal vector field u, (x) and marginal rate function (x) solve the dynamical WFR problem. The proof is completed.
A.6 PROOF OF PROPOSITION 4.3
Proposition 4.3 For balanced cases, let 5 + 30 (3.3) then the travelling Dirac
3.5
reduces to
uniform rectilinear motion with conserved mass, and the travelling Gaussian (4.)) simply reduces
to un (z|z) = Φ -0, (z(x(2) = O, thus we recover OT-CFM (Tong et al. 2024a
Proof. With conserved mass, and oo, the growth rate is forced be 0 in order make the cost finite. Thus, the minimizer of dynamical WFR (.3) simply reduces to the minimizer of dynamical
OT
inf /' / |u(x 1) (x: (d)dddl
WFR (Mo.
P.g.u Jo 2
(A.33)
s.t.
0.p + V (pu) = pg: = 0, Po = Ho: P1 = f1.
The corresponding OET 3.8 minimizer also reduces the static OT minimizer
22

---

Published as a conference paper at ICLR 2026
rEM, (22) 11 -2 In cos(
| 1-yll2 (r(ee drdd
WFR u1) 28²
inf
20
+ KLL / (x, y)dyyllolxxx)) + KL(
/ (z, y)d ()dx)|11(()))
) 1//8 -21n(1
||
Uاا +(00 ydddy
25
inf
1EM4(x3)
28
+ +KI. / 70 (ddy))/((())) + KT, / (2,
1/ (( x (1---yl yll2 )² + 0(8 ) ydaady
=22
inf
(A.34)
1EM
28
+ KLL /// 7(0,, yddyl (z)) + KL / (xx y)d
U/ee -11- y1 y)drdy + o(1)
inf
1(x, yydyylyllll))) + 25t KL, / 7(x, ()dx)((1(())
+ 28 KL( x
inf (T:7
x² 21 11, dxdd
The semi-coupling constructed from the OT coupling is just the OT coupling Thss, we have mi = mi for each pair of points sampled from the semi-coupling. The mass is denoted m in the following.
The cost between two dirac 3.4 reduces
2 112
WFR-DD (m0 mδσ1)= 282 m 2mcos(
28
|1:00-
221122)
4m8² (1 cos(
26
| xill2
4m8² (
j2
(0(5-2)
(A.35)
28
1
2 5m|11000
x x2112 + o(1)
8:200 1
mm| x 2
m
2
When 5 - co, all pairs of æo æ satisfies ||00 - 2 æ1112 < कτ. The travelling Dirac 3.53.6
reduces
to uniform rectilinear motion with conserved mass following
| - -0oll2
[ =zoll2 0 o(1)
tan
26
28
= 28T V
m²
1= (1|æ))
æ0l12 2 + o(1))
5,200 m||x
20
1 72
/1 + (1)
A 2m
2m
2m(1
1
6 6::00 0
v1 + 72
V1+ o(1)
(A.36)
B m
m(1
1
6:300 0
VI T
VI+ o(1)
(l))
A12
2Bt + m
80200
u(l)=
8 80200
| Bi 2|
By
ml)
Thus, the tavelling Gaussian
4.8 reduces to
po (0r Do, 21) A N poo ((x1 02)
23

---

Published as a conference paper at ICLR 2026
The conditional velocity field and growth rate reduces to
Ut (3/22)
at 20, gr [x(])) O2 In
0
The proof is completed.
A.7 PROOF OF PROPOSITION 5.1
The solution to the multi-time WFR problem equation 5.7 is equivalent to the concatenation of the solutions of successive time points.
Proposition 5.1
Proof. Let to < 41 <.. < Lx be the discrete time points. For each & 1,..., K, denote by (pl t), (2c, t),g g) (x,t) the optimal solution of the two-time WFR problem between 1oa-1 and pea:
tk
inf f
/ (Iu(ar, f)l2 80 p(xr, dx dt (A.37)
p.u. tu
subject to
dep + Vn (pr) = pg: Pie 1 = Mta. 15 Phu = pmu
(A.38)
Now define the concatenated trajectory
f( 1) (k) (ar, (), ū(x; 1) = (2) (x, 1), G(xr () =g (x, for [[=-1, (1))
(k+1) the concatenated triple (p. ü, g) is admissible for the multi-time WFR Since Pix = 1ll = Pie problem.
The corresponding action is
tk
to
f / (1)ū(z, 6 t)1l3) p(x; dx dt
lx
2 ta x
(A.39)
tk to
WFR (
th -
k
Hence
K lk to
3 })
WFR
(pt-s11:)))))
(A.40)
tv
t
To see that equality must hold, assume by contradiction that there exists another admissible solution (p', u,, such that
tk to
1.a rtx
1 (1) (a, o' dx dt
Σ
K
tk 1 to 4x-1
( Prx)
2
Since the total action sum over the disjoint intervals [1 [u], this strict inequality implies that there must exist at least one index k* such that
[K
to
/ (Ilu'(2, 1)12+ 1 (æ, o' (xc, 1) dar di < tK to tx- f---1
WFR?
(pt
Mexe
2
Dividing both sides by the positive factor
tg-to 4 4x-
gives
t
tv
Jg / (1lu (0,()))))))) 1)12 + (x, (ax, da di < WFR
2
But this contradicts the definition of WFR (Pt. 1- Me as the minimal action between (l.... and l4g.. Hence no admissible solution can have strictly smaller action than the concatenated one.
We conclude that the concatenated trajectory
K
(0' (z, = (R) (1)) (H)), tE (tt-1,,t tul,
(A.41)
1=1
24

---

Published as a conference paper at ICLR 2026
achieves the infimum of the multi-time problem and is therefore optimal.
Remark: Proposition 5.1 holds true when the transport between adjacent time pairs (th-1: tk) is independent. However, in marti-aarginal methods, such as MMFM (Rohbeck et al. 2025), 3MSBM (Theodoropoulos et al. 2025) and MMSFM (Lee et al. 2025). which enforce global coherence across all marginals, the dependencies between adjacent time pairs are introduced. As result, the assumption of independence no longer holds, and Proposition 5.1 does not apply in the same form under multi-marginal framework.
B ADDITIONAL RESULTS
B. EXPERIMENTAL DETAILS
The experiments were performed on a shared high-perffrmannce computing cluster with NVIDIA A100 GPU and 128 CPU cores. The architecture of the neural networks used to parameterize vo(x, t) and Iø (x. t) are Multilayer Perceptrons with 256 hidden channels and layers. We use the LeakyReLU activation function. These networks were optimized using Pytorch (Paszke et al.) 2017). The OET problem solved using the Python Optimal Transport (POT) package (Flamary al. 202T
B.2 EVALUATION METRICS
We use two metrics to evaluate model performance: the I-Wasserstein distance (W)) to measure the similarity between normalized predicted and true distributions, and the Relative Mass Error (RME) to assess how well the model captures cell population growth.
The metries are defined as:
aEll(p.)) / ||x (1127(((1 8)
(i(p.))
min
Ex we((s) nx/nol
RME(tx)
nk/no
Here, p and q are the predicted and true empirical distributions, wx(tx) is the inferred weight of cell i at time l. and is the number of cells at time k.
To evaluate a model, we apply its learned dynamics to initial cell data (where all weights are uniform, (0) 1/mg) to predict cell trajectories and their corresponding weights at subsequent time points if the unbalanced term is included, otherwise the weights are kept uniform. We then compute the weighted Wi and RME by comparing these predictions to the observed cells from dataset. For some datasets, we also perform a hold-out experiment, training on all but one time point and using W1 to evaluate performance the performance of TIGON (Sha et al.) 2024), we reimplemented their method to avoid numerical instabilities. For other methods, we mainly used their default settings if the datasets included were used in their original paper, otherwise we changed the network sizes of these models to ensure comparable parameter counts, tuned the training epochs and learning rates for these models suit different datasets for fair comparisons. Specifically, for VGFM, as most datasets used in our work are included their original codebase, we directly adopted their hyperparameters. For other simulation- free methods, we adjusted the training budget based on data complexity. For high-dimensional real- world datasets (such as EB, CITE, Mouse), we increased the training duration to 3,000 epochs and incorporated learning rate schedulers (Cosine Annealing) to ensure stability and convergence. For simulation-barl methods (MIOFlow, TIGON and DeepRUOT), we adopted a distinct but unified hyperparameter configuration. To ensure consistent and fair comparison among these methods, set their hyperparameters to the values recommended by Deep which represents the current hate-of-the-ar in this category.
the unseen time point. Specifically, to evaluate
25

---

Published as a conference paper at ICLR 2026
B.3 SIMULATION GENE DATA
We adopt the same dataset in Zhang et al. Its dynamics governed by the following set of stochastic ordinary differential equations:
2025a) that simulates a synthetic gene regulatory network.
dX1
x1 3
61 + 7/1
dt
+
+ 72 13 +
dX2
02.23 B
82 + 725
1 1+
(2
73X +
dX3
a3
d3X3
di
03Xg
Here, (t) is the concentration of gene i. The model features a toggle switch between Xı and X2 (mutual inhibition and self-activation), which are further inhibited by Xa and activated by an external signal B. The equations are parameterized by rates for self- activation (a,), inhibition (x), and degradation (6,), with an added stochastic noise term (7)E)).
The simulation included probabilistic cell division, where the instantaneous growth rate, g. depends x3 on the expression of X2 ! og 1+x;- Upon division, daughter cells inherit the parent's state plus small perturbation. Data was recorded time points [0, 8, 16, 24, 32] from distinct initial cell populations. One initial population exhibits transition and growth, while another population is steady state. shown Figure 4 WFR-FM able to recover the increasing growth rate of the population the right side, eliminating false transition.
a
an
3.40
4M
Figure 4: (a) True growth rate (b) Predicted growth rate by WFR-FM
Choice of Growth Penalty The parameter å represents the penalty of growth our algorithm. An excessively low 8 would incorrectly incentivize solutions where cells tend to vanish and appear again to match the observed distributions, which may not be biologically plausible. In contrary, an overly high 8 would excessively suppress cell proliferation. Therefore, reasonable choice of is to balance transition and growth, which requires to scale appropriately with the data. To evaluate the effect of 8, further conducted ablation study the simulation gene dataset. As shown in Table 5) both an overly high (6=3.0) and low (8=0.5) value result in sub-optimal performance. However, the results remain comparable for d within suitable range (1.0 2.0), demonstrating the robustness of our algorithm with respect to this parameter, provided it matches the scale of data. set to 1.5 for the results of this dataset in the main text.
B.4 DYGEN DATA
We adopt the same dataset as in Huguet et al. (2022); Wang et al. (2025) which was simulated by Dygen (Cannoodt al. 2021) model unbalanced bifurcation. The dataset contains 728 cells, and its dimension was reduced to 5 using PHATE (Moon et al.) (019). The data is challenging as it exhibits varying mass across different time points, and the number cells in the lower branch larger than the upper branch. As shown in Figure 5 WFR-FM accurately generates the bifurcating trajectories, and predicts higher growth rate in the lower branch, capturing the unbalancedness the bifurcating process with the growth penalty set 2. We also summarized the detailed quantitative results across different time points in Table As shown, WFR-FM achieves the lowest Wi, and RME in most of the time points.
26

---

Published as a conference paper at ICLR 2026
Table 5: Sensitivity analysis for parameter å on simulation gene dataset.
Parameter
t=1
t=2
t=3
t=4
W1
RME
W1
RME
W1
RME
W1
RME
0.5
0.021
0.003
0.036
0.016
0.036
0.022
0.033
0.027
1.0
0.020
0.001
0.021
0.002
0.018
0.003
0.015
0.001
1.5
0.021
0.001
0.021
0.001
0.018
0.002
0.016
0.001
0.022
0.000
0.025
0.000
0.025
0.003
0.032
0.000
2.5
0,023
0.001
0.042
0.003
0.025
0.006
0.042
0.002
3.0
0.027
0.001
0.044
0.003
0.051
0.006
0.030
0.004
Cante-
Gound Praccery
Figure 5: (a) Learned trajectories (b) growth rate the dygen dataset
Table 6: Comparison of method performance over time on the Dygen dataset. Best results are in bold.
Method
t=1
t=2
t=3
t=4
Wi
RME
W1
RME
W1
RME
W1
RME
MMFM
0.574
1.704
1.499
1.706
Metric FM
0.892
2.347
2.030
1.799
SF2M
0.637
1.266
1.415
1.790
MIOFlow
0.420
0.640
1.537
1.263
TIGON
0.446
0.033
0.584
0.060
0.415
0.023
0.603
0.071
DeepRUOT
0.454
0.011
0.481
0.070
0.870
0.104
0.688
0.074
Rar-RUOT
0.315
0.128
0.548
0.336
0.630
0.222
0.593
0.023
UOT-FM
0.652
0.008
0.780
0.077
1.252
0.090
2.130
0.213
VGFM
0.335
0.001
0.312
0.073
1.109
0.041
0.634
0.033
WFR-FM
0.110
0.003
0.098
0.007
0.211
0.008
0.121
0.002
27

---

Published as a conference paper at ICLR 2026
B.5 GAUSSIAN MIXTURE DATA
We adopt the 1000D Gaussian Mixture Data from Wang et al. (2025). Following their setup, an initial distribution of 500 samples (100 from upper Gaussian, 400 from lower Gaussian) and final distribution of 1,400 samples (1,000 from the upper, 200 from each lower two Gaussians) were generated. This models cell proliferation in the upper region, which serves as a suitable example to evaluate the algorithm's 's ability of modeling unbalanced dynamics in high dimensional settings. As shown in Figure 6 WFR-FM correctly models the higher growth rate in the upper region, and yields correct trajectories with growth penalty as 1.4, showing its effectiveness.
b
0
Cone
Figure 6: (a) Learned trajectories (b) growth rate on the Gaussian 1000D dataset
B.6 EPITHELIAL MESHHIHHMAA TRANSITION DATA
We adopt the dataset that captures the poattheiallmessnhhymne transition (EMT) in A549 lung cancer cells from Cook& & Vanderhyden 2020). The dataset includes samples collected four distinct time points throughout this process. The dimension was reduced to 10 using an AutoEncoder (AE) as in Sha et al. (2024). As shown in Table 7 WFR-FM achieves the best performance in both distribution matching task and mass matching task across all time points with growth penalty set to 2. We plotted the learned trajectory and growth rate in Figure 7 As shown, WFR-FM predicts higher growth rate in initial and intermediate stages, which is in line with the results predicted by DeepRUOT (Zhang al. 2025a) these cells exhibit enhanced stemness.
Table Comparison of method performance over time the 10D EMT dataset. Best results are in bold.
Method
t=1
t=2
t=3
Wi
RME
Wi
RME
Wi
RME
MMFM
0.2576
0.2874
0.3102
Metric FM
0.2605
0.2971
0.3050
SF2M
0.2566
0.2811
0.2900
MIOFlow
0.2439
0,2665
0.2841
TIGON
0,2433
0.002
0.2661
0.003
0.2847
0.001
DeepRUOT
0,2902
0.001
0,3193
0.011
0.3291
0.002
Var-RUOT
0.2540
0.075
0.2670
0.014
0.2683
0.041
UOT-FM
0.2538
0.002
0.2696
0.013
0.2771
0.010
VGFM
0.2350
0.016
0.2420
0.011
0.2450
0.018
WFR-FM
0.2099
0.001
0.2272
0.002
0.2346
0.001
B.7 EMBRYOID BODIES DATA
We adopt the dataset of 16,819 cells from a 27-day differentiation time course of human embryoid bodies (EB), sampled at five time points (Moon et al. 2019). This experimental system models early embryonic development. The original high-dimensional gene expression data was preprocessed by reducing it to lower dimensions using Principal Component Analysis (PCA) as in Wang al. (2025)
28

---

Published as a conference paper at ICLR 2026
GereA,
Figure 7: (a) Learned trajectories (b) growth rate on the EMT dataset
Effects Mini-batch WFR-OET on WFR-FM. We employed the mini-batch WFR-OET to ac- celerate computation on this dataset. The effects of different batch sizes are summarized in Table 8 The mini-batch approach significantly reduces runtime compared to the full-batch method. we noticed that the runtime exhibits non-monotonic trend, which may due to the higher number of OET calculations required for smaller batches. Performance generally improves as the batch size increases, eventually becoming comparable to the full-batch result. Therefore, selected the batch size of 2000 as the default value to achieve a favorable balance between efficiency and accuracy.
Table 8: Sensitivity analysis for batch size of mini-batch WFR-OET on the 100D EB dataset.
Batch Size
t=1
t=2
t=3
t=4
Time (s)
Wi
RME
Wi
RME
W/
RME
Wi
RME
500
9.980
0.010
11.075
0.008
11.567
0.008
12.850
0.004
159
1000
9.956
0.009
11.076
0.008
11.609
0.006
12.783
0.002
129
1500
9.937
0.009
11.046
0.008
11.583
0.005
12.771
0.001
125
2000
9.941
0.009
11.040
0.006
11.516
0.008
12.664
0.005
131
3000
9.914
0.009
11.049
0.009
11.529
0.006
12.687
0.002
139
4000
9.920
0.009
11.050
0.007
11.538
0.006
12.654
0.003
159
w/o mini-batch
9.926
0.009
11.045
0.007
11.545
0.006
12.657
0.002
269
Effects of Data Dimensions on WFR-FM. To evaluate the scability WFR-FM, we compare the performance of WFR-FM with other methods on 5 dimensions (Table IO with growth penalty set to 2, 50 dimensions (Table TT) with growth penalty set to 25 and 100 dimensions (Table [2) with growth penalty
35. WFR-FM outperforms other methods in in most of the time points all three datasets, indicating the consistent performance of WFR-FM when scaling to higher dimensions with little increase in computational time (Table 9
B.8 CITE-SEQ DATA
We adopt the CITE-seq dataset from Lance et al. (2022) which contains 31 31,240 cells collected four time points. Following the preprocess procedure in Wang et al. (2025), we selected only the gene expression matrix and projected it into 50 dimensions using PCA. We summarized the performance of distribution matching and mass matching in Table [3] We adopt the mini-batch strategy with the batch size of 2000 for this dataset, and set growth penalty to 30. shown, WFR-FM achieves
Table 9; Computational time across different dimensions the EB dataset.
Dimension
Time (s)
5
123
50
126
100
131
29

---

Published as a conference paper at ICLR 2026
Table 10: Comparison of method performance over time on the 5D EB dataset. Best results are in bold.
Method
t=1
t=2
t=3
1=4
Wi
RME
Wi
RME
W1
RME
W1
RME
MMFM
0.477
0.554
0.781
0,872
Metric FM
0.449
0.552
0.583
0.597
SF2M
0.556
0.715
0.750
0.650
MIOFlow
0.442
0.585
0.651
0.670
TIGON
0.386
0.002
0.502
0.015
0.602
0.021
0.600
0.027
DeepRUOT
0.386
0.005
0.497
0.017
0.591
0.021
0.585
0.030
Var-RUOT
0.416
0.111
0.486
0.144
0.509
0.054
0.511
0.022
UOT-FM
0.544
0.032
0.670
0.029
0.729
0.016
0.852
0.041
VGFM
0.402
0.046
0.494
0.018
0.525
0.035
0.573
0.021
WFR-FM
0.324
0.003
0.401
0.001
0.431
0.005
0.510
0.005
Table 11: Comparison of method performance over time on the 50D EB dataset. Best results are in bold.
Method
t=1
t=2
t=3
t
W
RME
WI
RME
W1
RME
Wi
RME
MMFM
9.124
10.474
11.022
11.480
Metric FM
8.506
9.795
10.621
12.042
SF2M
9.247
10.882
11.650
12.154
MIOFlow
8,447
9,229
9.436
10.123
TIGON
8,433
0.067
9.275
0.022
9.802
0.179
10.148
0.101
DeepRUOT
8.169
0.003
9.049
0.038
9.378
0.088
9.733
0.004
Var-RUOT
9,442
0.128
9.709
0.081
10.482
0.031
10.735
0.030
UOT-FM
8.717
0.063
10.858
0.009
11.813
0.022
12.733
0.018
VGFM
7.951
0.089
8.747
0.042
9.244
0.019
9.620
0.044
WFR-FM 7.664
0.008
8.659
0.006
9.182
0.004
9.914
0.004
highest accuracy in matching distributions across all time points, while maintaining a mass matching accuracy comparable other method that model unbalancedness.
B.9 MOUSE HEMATOPOIESIS DATA
We adopt the mouse blood hematopoiesis dataset (Mouse) from Weinreb et al. (2200). We selected 49,302 cells with lineage tracing information collected at three time points. The original gene ex- pression data was reduced to 50 dimensions using PCA. This dataset was chosen for its substantial cell population growth, making it suitable for scalability testing. We sub-sampled the dataset to
time 000
135
75
Number Coll
Figure 8: Scaling of runtime with respect cell number
30

---

Published as a conference paper at ICLR 2026
Table 12: Comparison of method performance over time on the 100D EB dataset. Best results are in bold.
Method
t=1
t=2
t=3
1=4
W1
RME
W1
RME
Wi
RME
W1
RME
MMFM
11.460
13.879
14.441
14.907
Metric FM
10.806
12.348
13.622
16.801
SF2M
11.333
12.982
13.718
14.945
MIOFlow
11.387
12.331
11.905
12.908
TIGON
10.547
0.014
12.926
0.052
13.897
0.107
14.945
0.096
DeepRUOT
10.256
0.002
11.103
0.074
11.529
0.136
12.406
0.047
Var-RUOT
11.746
0.091
12.237
0.024
12.957
0.150
13.335
0.074
UOT-FM
10.757
0.056
12.799
0.037
13.761
0.044
15.657
0.022
VGFM
10.313
0.048
11.278
0.035
11.703
0.028
12.637
0.066
WFR-FM
9.941
0.009
11.040
0.006
11.516
0.008
12.664
0.005
Table 13: Comparison method performance over time on the 50D CITE dataset. Best results are in bold.
Method
t=1
t=2
t=3
Wi
RME
Wi
RME
W
RME
MMFM
33.971
36.854
43.721
Metric FM
28.314
28.617
33.212
SF2M
29.543
32.655
36.265
MIOFlow
28.290
28.524
32.230
TIGON
28.196
0.186
27.921
0.545
32.846
0.653
DeepRUOT
28.245
0.168
27.908
0.525
32.950
0.634
Var-RUOT
30.219
0.331
32.702
0.325
40.613
0.486
UOT-FM
33.531
0.009
32.795
0.046
49.751
0.097
VGFM
29.449
0.020
29.722
0.057
33.752
0.001
WFR-FM
27.831
0.043
27.478
0.045
34.784
0.022
different cell numbers, and plotted the scaling of time with respect to cell numbers in Figure 8 As shown, the training time scales linearly with cell numbers, in line with other competing meth- ods, indicating that WFR-FM with the batch size 2000 for this dataset, and set growth penalty to 15. As shown in Table [1] WFR-FM outperforms other methods by large margin, indicating the effectiveness of WFR-FM in simultaneously modeling cell proliferation and differentiation.
be applied larger datasets. We adopt the mini-batch strategy
1.10 UNBALANCEDNESS IN REAL SINGLE-CELL DATASETS
We illustrate the proportions of cell types (or clusters) across different time points in several real single-cell RNA-seq datasets (Moon et al. 2019 Lance et al. 2022 Weinreb al. 2020). As shown in Figure 9 the relative abundances of different cell types vary substantially across time points, indicating a pronounced unbalancedness in these datasets. We hypothesize that this widespread imbalance may partly explain why methods based on unbalanced optimal transport better than those assuming balanced distributions. Furthermore, previous study (Sha et al.) (024) has also reported the impact of cell proliferation on dynamics in the EMT (Cook & Vanderhyden, 2020) and Mouse (Weinreb et al.] 2020) datasets, highlighting the importance of accounting for growth and apoptosis in trajectory inference.
perform
31

---

Published as a conference paper at ICLR 2026
Table 14: Comparison of method performance over time on the 50D Mouse dataset. Best results are in bold.
Method
t=1
t=2
W1
RME
W1
RME
MMFM
7.647
10.156
Metric FM
7.788
11.449
SF2M
8.217
11.086
MIOFlow
6.313
6.746
TIGON
6.140
0.382
6.973
0.326
DeepRUOT
6.052
0.062
6.757
0.041
Var-RUOT
7.951
0.131
10.862
0.154
UOT-FM
8.114
0.035
9.170
0.011
VGFM
6.274
0.076
6.796
0.070
WFR-FM
5.486
0.012
6.211
0.011
EB
CITE
Call bype/ Cluster proporsion
1.0
0.8
0.8
EryP
10
Ouuserr
0.6
MesP
0.4
12
MoP
13
0.2
0.2
0.0
0.0
Time
Time
Mouse
Coll bypt/ Cluser propertion
1.0
Baso
Cor7,DC
0.8
Erythraid
Lymphaid
0.4
Most
0.2
Monocyte
Neutrophil
Undifferentiated
Time
Figure 9: Proportions of cell types (or clusters) across different time points for three real scRNA- seq datasets. Top left: EB dataset Moon et al. (2019); top right: CITE dataset Lance et al. (202)):
bottom left: Mouse dataset Weinreb et (2020).
32

---

Published as a conference paper at ICLR 2026
C RELATION TO OTHER UNBALANCED FLOW MATCHING METHODS
C.I RELATION TO UOT-FM.
UOT-FM (Eyring et al. 2024) was proposed to extend flow matching the unbalanced setting. Its formulation replaces the balanced OT map in OT-FM (Tong et al. 2024a) with an unbalanced OT map obtained from the classical static unbalanced OT problem:
(23) E/ee c(ar, vn, y)dady
UOT m)) inn
(C.1)
+ + KL( / 7<x y)dy)llo(zz) + X2 KL( / 7(, y y)dzllu
This provides a natural way relax mass conservation flow matching by leveraging static un- balanced OT formulation. In comparison, WFR-FM is built on the dynamic WFR geometry, which directly yields a continuous-time interpretation of the unbalanced dynamics. In addition, while UOT-FM focuses on learning the transport vector field, our formulation simultaneously learns both a displacement field and a growth rate function, enabling us to model single-cell dynamics where mass and state evolve jointly over time.
C.2 RELATION TO UFM.
UFM (Corso et al. 2025) can be viewed as a natural generalization of UOT-FM. Rather than restrict- ing the coupling the optimal plan obtained from static unbalanced OT, UFM allows for arbitrary couplings between source and target samples. In particular, UFM focuses on arbitrary couplings and provides both a formal and empirical characterization of the associated objectives. The training objective augments vector field regression with additional terms that penalize mass mismatch:
LUPM(7,0) or (xE,, [[xx(();)) - (x (Xo. x ))
(C.2)
+ A1 KL( / r(x y)dylluo())) + X2 KL(/ / 7(x, y)d (y))
This formulation extends unbalanced flow matching beyond the static UOT setting by simultane- ously learning transport dynamics and accounting for marginal discrepancies. Compared to this, WFR-FM is derived from the dynamic WFR geometry, which provides principled interpretation of both displacement and growth within a unified continuous-timg framework.
C.3 RELATION TO CAO ET AL. 2025
Cao al. 2025) ) adopt a formulation that conceptually similar to UOT-FM, but specializes it the concrete application of pansharpening. On top of the basic unbalanced flow matching framework, they introduce additional design choices tailored to the data characteristics of this imaging task. While effective for pansharpening, such a task-specific adaptation faces challenges in generalizing to the reconstruction of single-cell dynamics.
C.4 RELATION TO VGFM
VGFM WWang et 2025) is related work ours as both approaches extend OT-based flow matching to the unbalanced setting by jointly learning a transport vector field and a growth rate. VGFM built upon a conitconstra OT formulation:
SC-UOT m1) = inf / ||x (x1 + +λ1 KT. / 7(,,, v)dvllmo())) 720
x
(C.3)
[ (x y)dz p1()),
subject to
This formulation leads to dynamic interpretation where mass variation and transport sepa- rated into two distinct phases first birth- death dynamics without transport, followed by transport without mass change. For single-cell applications, this corresponds to cells proliferating dying
33

---

Published as a conference paper at ICLR 2026
without altering gene expression in one phase, and changing gene expression without proliferation or apoptosis the other. However, such separation remains biologically challenging to justify. To address this, VGFM introduces additional heuristics couple transport and growth; however, this modification moves the method away from a strict OT formulation. Furthermore, VGFM is primar- ily employed as pretraining strategy and still requires simulating ODE trajectories (albeit fewer than without pretraining) to achieve strong performance.
By contrast, WFR-FM is derived directly from the dynamic WFR metric, couples transport and growth within a principled geometric framework, and achieves fully simulation-free training without any additional ODE integration.
D INTUITION AND LIMITATIONS OF MINI-BATCH OT
In WFR-FM, we adopt mini-batch approach approximate the optimal transport coupling during training. The convergence of mini-batch optimal transport to the true coupling remains an open problem. For the balanced case, some theoretical and numerical studies have explored how batch size and sample number affect convergence (e.g.. Sommerfeld et al. (2019)), however, to our knowl- edge, analogous results for the unbalanced setting are not yet available. Intuitively, one can view an independent coupling as the extreme case of mini-batch OT with batch size I, while full-batch OT corresponds to the other extreme. Selecting an appropriate batch size thus reflects trade-off' between computational efficiency and adherence to the optimal transport principle. Importantly, even with mini-batch OT, our method remains a flo-matcciig approach: continues to match the marginals the data distributions. Although mini-batch OT may result in slight deviations from the globally optimal transport plan, it provides a practical and effective way to learn the underlying transport and growth dynamics.
In addition, we conducted an ablation study on the choice of OT batch size (see Table 8). The results show that our method is robust to the selection of batch size: performance remains stable across wide range of batch sizes, indicating that mini-batch OT does not introduce significant sensitivity in practice.
E PSEUDOCODE FOR INFERENCE IN WFR-FM
Algorithm 2 WFR-FM Inference Workflow
Require: Initial samples {2r I.V Mr. trained vector field vo (), trained growth rate 9o (2, t), target time fend, integration step size At.
1: Initialize a:()) +- (0) m(()) for I.....N
2: to
3: while lend do
4: Atscep f- min At, tend t)
ensure last step reaches exactly food
5: for = - N do
6:
vo (1)
7:
g (=(()())
Update position: (0() (+) WtSepp Update mass: (i) mi) exp(gAtuup
8:
9:
10:
end for
11: 1+1+ Alsep
12:
end while
return Evolved samples (zx(1), ml()) J
F THE USE OF LARGE LANGUAGE MODELS (LLMs)
We used L.LMs solely for the purpose of improving grammar and clarity of the text in manuscript.
34


Reviews:
[{"review_id": "vf1q8MEcfY", "reviewer": "Reviewer_XcwL", "summary": "This paper proposed use of Wasserstein-Fisher-Rao (WFR) metric to capture unbalanced transport by jointly regressing both a transport vector field and growth function, while explicitly modeling growth dynamics unlike other unbalanced transport FM methods. The method is evaluated on a range of synthetic datasets as well as commonly used real-world datasets for single-cell biology application.", "questions": "* In line 115 and 118, authors cite work by Neklyudov (2023, 2024) [1,2] and Sun (2025) [3] providing commentary that existing unbalanced approaches lack framework that can jointly recover velocity and growth dynamics or require costly ODE simulations. Would it be possible to evaluate the quality of learnt growth dynamics in a synthetic setting and compare to other unbalanced transport baselines such as Neklyudov (2023, 2024) and Sun (2025)?\n* Proposition 5.1 states that the solution to the multi-time WFR problem is equivalent to the concatenation of the solutions of consecutive time points. From algorithm 1 it also seems that the shared time-continuous velocity network is trained on consecutive pairs of points? However recent multi-marginal methods such as 3MSBM by Theodoropoulos et al (2025) [4] jointly enforce all marginals achieving better global coherence. Does proposition 5.1 then still hold?\n\n**References**\n\n[1] Neklyudov, Kirill, et al. \"Action matching: Learning stochastic dynamics from samples.\" International conference on machine learning. PMLR, 2023.\n\n[2] Neklyudov, Kirill, et al. \"A computational framework for solving wasserstein lagrangian flows.\" arXiv preprint arXiv:2310.10649 (2023).\n\n[3] Sun, Yuhao, et al. \"Variational Regularized Unbalanced Optimal Transport: Single Network, Least Action.\" arXiv preprint arXiv:2505.11823 (2025).\n\n[4] Theodoropoulos, Panagiotis, et al. \"Momentum Multi-Marginal Schr\\\" odinger Bridge Matching.\" arXiv preprint arXiv:2506.10168 (2025).", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "* **Clarity**: The authors provide comprehensive theoretical background in main text and appendix as well as relation to other unbalanced OT related work highlighting key differences between baselines and proposed method. The authors further support theoretical background with extensive proofs in the appendix.\n* **Robustness**: Robustness is tested through several ablation studies on batch size and coupling as well as controlled synthetic settings.\n* **Low computational cost**: The authors suggest method that improves performance in matching end-point marginals in comparison to baselines while keeping low computational cost. Further authors also show that the computational cost does not significantly increase across different dimensions or batch sizes.", "weaknesses": "* **Evaluation of learned growth dynamics**: The paper demonstrates superior performance across $W_1$ and $RME$ metrics, however paper would benefit from showing it also manages to learn correct growth dynamics given it explicitly learns term $g_{\\phi}(x,t)$. This should be added as Q5 in experiments section and validated in synthetic or real-world setting.\n\n* **Adding further overview of use of WFR for FM-based algorithms**: It would be good to see more extensive discussion and if possible empirical comparison to unbalanced transport action matching (AM) by Neklyudov (2023) which also uses WFR distance to construct unbalanced OT. It would be good to see trade-off between using WFR to explicitly learn growth dynamics and using a single network to learn both transport and growth dynamics as presented in AM", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mMhsaz1tU7", "reviewer": "Reviewer_V2He", "summary": "The paper introduces WFR-FM, a model for learning continuous cellular dynamics from single-cell snapshot data when population sizes change over time. WFR-FM is based on the Wasserstein–Fisher–Rao (WFR) geometry, which couples transport of cell states with birth–death mass changes. To learn these dynamics efficiently, the authors combine WFR with flow matching, producing a simulation-free training objective. WFR-FM learns two components of cell dynamics jointly: a velocity field that transports cell states through gene expression space, and a growth rate field that models local changes in cell population density. The method constructs conditional particle trajectories inspired by traveling WFR geodesics, and trains a neural network to regress both velocity and growth directly, without integrating differential equations during training. The primary use of WFR-FM is to infer trajectories from time-series scRNA-seq data, particularly when cell counts shift between time points. The authors show that WFR-FM reconstructs developmental and differentiation dynamics more accurately and efficiently than ODE-based dynamic optimal transport methods, while capturing biologically meaningful patterns of proliferation and apoptosis.", "questions": "- L 308: Could you provide more intuition on the unit velocity here and why it is equivalent to flow matching by Tong?\n- Did the authors inspect any theoretical guarantees of their batchwise approach (just asking for an intuition)? How detrimental is it to use this approach for learning the true coupling?\n- Why would unbalanced OT work better on the EMT, EB, and cite datasets? I expect it to work on par with other OT and FM models, but why would it overcome them? Is it related to some aspect of umbalancedness?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Overall, I enjoyed the paper. I think it is clean, well-written, and the math flows quite well. Broadly, I also find the research timely and compelling. While existing publications have already studied how to approach unbalancedness in flow matching, the present paper provides quite a deep formulation of the problem with updated probability paths and proper justifications. The model also seems to be working well in applied scenarios.", "weaknesses": "Overall, my evaluation of the paper is positive (see the score). However, some aspects where I would suggest some improvements. I will elaborate on them below. \n\n- **Minor:** L136 -- change the citation to `\\citep` for Tong et al. \n- In Eq. 3.3, $\\delta$ is undefined. Since later the same symbol is used to define Dirac measures, I think it would be better to differentiate the meaning of $\\delta$ in different scenarios. \n- **Consistency in velocity notation:** In 3.5, the velocity notation is $x'$, but this is not used anywhere else in the text, where velocity is mostly expressed in terms of differentials. \n- L195: \"They\" does not seem like a correct pronoun here. What's the subject?\n- Eq. 3.8: The OET misses the constraints.\n- In Eq. 4.7: Refer to $\\tau$ as the $\\tau$ described above. \n- The way the traveling Gaussian is introduced could benefit from more elaboration. Especially, you could add one sentence explaining why the mass $m_0=1$ (hence, cause it's a relative mass change along the flow). \n- The concept figure could be bigger with a more exhaustive caption.\n- In the paper, I would probably describe the RME metric with a few sentences, and then refer to the appendix. At the moment, there is no clear intuition about it in the main text, and this breaks the readability of the experimental section a bit. Also, acronyms are not expanded in the main. Similarly, I would add a few sentences about where the presence of unbalancedness plays a role in the chosen datasets and why accounting for it is expected to produce a better performance. \n- Is the action in Tab. 2 calculated over multiple trajectories? If yes, can you please add error bars? Along these lines, I think that most of the tables would benefit from error bars, as the results are very close.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jpjVztags2", "reviewer": "Reviewer_Tvju", "summary": "This paper tackles the important and challenging problem of modeling continuous dynamics from sparse, unbalanced temporal snapshots --- a problem of significant interest in single-cell biology, where cell proliferation and apoptosis are key factors. The authors propose WFR Flow Matching (WFR-FM), a novel, simulation-free framework that unifies dynamic unbalanced optimal transport (under the WFR geometry) with flow matching.\n\nThe core technical contribution is a method that, unlike standard flow matching, simultaneously regresses both a transport vector field for particle displacement AND a scalar growth rate function for birth-death dynamics. The authors provide theoretical guarantees that this approach exactly recovers WFR geodesics. Empirically, the paper demonstrates that WFR-FM achieves state-of-the-art performance on several single-cell trajectory inference benchmarks, outperforming a wide range of baselines in accuracy, efficiency, and stability.", "questions": "1. **Practicality of Weighted Samples**: Regarding the practical application: As I understand from the inference workflow (Algorithm 2), the number of simulated particles/samples remains fixed from $t_0$. The 'growth' is modeled by increasing the mass $m^{(i)}$ of each particle. Is this interpretation correct? If so, the final output is a set of weighted points, where the number of points equals the initial number. How does this weighted-point representation practically benefit downstream analysis in single-cell biology, beyond simply matching benchmark distribution measures (like $\\mathcal{W}_1$ and RME)?\n\n2. **Coupling of Growth and Transport**: In the governing ODEs (Eq. 3.1), the evolution of the position $\\phi_t$ (driven by $u_t$) and the mass $m_t$ (driven by $g_t$) appear to be two separate, parallel processes, evaluated along the same path. Does the learned growth term $g_t$ (or the WFR-OET coupling) implicitly change the learned transport field $u_t$ for a given sample, compared to a model (like UOT-FM) that uses a similar unbalanced OT coupling but only models the vector field $u_t$?\n\n3. **Baseline Tuning**: This question refers to 2nd weakness.  Could the authors elaborate on the hyperparameter tuning procedure used for the baseline methods? Specifically, what steps were taken to ensure a fair and robust comparison, and to ensure that each baseline was performing optimally on these specific datasets?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. **Clarity**: The paper is well-written, clear, and easy to follow. The problem, the proposed method, and its theoretical underpinnings are all explained well.\n2. **Thorough Related Work**: The authors provide a comprehensive review of related work with broader context of optimal transport, flow matching, and single-cell trajectory inference.\n3. **Empirical Validation**: The experimental evaluation is a significant strength. The authors compare WFR-FM against a wide and appropriate range of state-of-the-art baselines (including both balanced and unbalanced OT/FM methods).\n4. **Scalability and Efficiency Analysis**: The inclusion of detailed scalability analysis, including training time and memory usage (e.g., Figure 3), is good.\n5. **Straightforward generalisation to multiple time points**: The method's design, which naturally extends to multiple time points by concatenating successive intervals (Proposition 5.1), is practical and well-justified for real-world snapshot data.", "weaknesses": "1.  **Positioning wrt to VGFM**: The paper currently relegates the detailed comparison with VGFM to Appendix C.4. While this appendix clearly articulates the important differences, VGFM appears to be the most direct competitor that also jointly models velocity and growth (and as reviewer believe was the first extension of flow matching to that). For fairness and clarity, a more prominent discussion of both the similarities and differences in the main body would be beneficial. This would also help readers better attribute the significant performance gains (e.g., in Table 1) to the novel methodological contributions of WFR-FM rather than potentially superior hyperparameter tuning\n2.  **Baseline Tuning**: The experimental results show some performance variations among baseline methods (e.g., MFM/SF2M vs. MIOFlow) that appear inconsistent across different experiments. The text lacks a detailed description of the hyperparameter tuning procedure used for these baselines. Without this, it is difficult to ascertain whether the reported baseline results represent their optimal performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)