TracingFlow: A Simulation-Free Trajectory Inference Framework Based on Second-Order Dynamics

TracingFlow: A Simulation-Free Trajectory Inference Framework Based on Second-Order Dynamics Yuhao Sun1∗, Zekun Wu2∗, Zixun Huang2 , Peijie Zhou1,3,4,5 † 1 Center for Machine Learning Research, Peking University 2 School of Mathematical Sciences, Peking University 3 Center for Quantitative Biology, Peking University

arXiv:2608.21070v1 [cs.LG] 21 Aug 2026

4 National Engineering Laboratory for Big Data Analysis and Applications, Beijing 5 AI for Science Institute, Beijing

Abstract Inferring continuous system evolution from sparse temporal snapshots is a key challenge in generative modeling and single-cell omics. While Optimal Transport (OT) is popular, existing frameworks are largely restricted to first-order dynamics, assuming memoryless velocity fields. This limits expressiveness, as first-order systems fail to account for regulatory momentum and time-delayed responses inherent in processes like cell differentiation. Here, we introduce TracingFlow, a simulation-free Flow Matching framework generalizing to secondorder dynamics. By using neural networks to regress the acceleration field, TracingFlow provides an exact, efficient solution to the Dynamical Optimal Acceleration Transport (DOAT) problem. Unlike first-order methods yielding over-smoothed trajectories, our second-order formulation captures high-curvature transitions and nonlinear evolutions by learning the underlying force fields. Evaluated on complex synthetic and large-scale scRNA-seq datasets, TracingFlow achieves superior accuracy in distributional reconstruction and trajectory faithfulness. Moreover, by integrating lineage tracing priors, it recovers dynamical structures that are both mathematically optimal and biologically plausible.

1

Introduction

Recovering underlying dynamics from discrete observations is a pivotal task in single-cell omics, known as Trajectory Inference (TI) [9, 72]. To identify the least costly evolution between distributions, Optimal Transport (OT) frameworks have emerged [7, 71]. These generally categorize into Static OT methods [22, 28, 51], which learn direct mappings, and Dynamical OT methods [24, 44, 57, 67], which capture continuous evolution via flow maps. Early Dynamical OT frameworks relied on Neural ODEs [9], incurring high computational costs due to numerical integration. To address this, Flow Matching [35, 36, 48] was proposed as a simulation-free alternative, regressing velocity fields to push source distributions to targets. By decomposing transport costs into single-particle trajectories, Flow Matching efficiently solves Dynamical OT problems [27, 50, 58], including its various variants [8, 15, 45, 59] . However, most TI frameworks assume first-order dynamics (ẋ = v(x, t)), inherently restricting v to be single-valued w.r.t x. This limits expressiveness when modeling complex biological priors, such as lineage tracing [30, 38], potentially yielding dynamics that contradict with additional biological prior information. Furthermore, [19, 23, 34] pointed out that if more modalities ∗ †

Equal contribution Corresponding author: pjzhou@pku.edu.cn

Preprint. August 21, 2026.

1

Figure 1: An illustration figure for TracingFlow such as chromatin accessibility or proteomics are considered, the biological kinetics may follow higher-order dynamics or could be modelled in augmented space. While 3MSBM [56] recently introduced a second-order Momentum Schrödinger Bridge for smoother trajectories, it relies on iterative retraining the acceleration field similar to Rectified Flow [36] to determine optimal couplings. Furthermore, it employs a heuristic method to estimate the initial velocity, rather than incorporating this estimation as part of the optimal control problem. To overcome these limitations, we introduce TracingFlow, a simulation-free framework utilizing second-order dynamical systems. We propose the Dynamical Optimal Acceleration Transport (DOAT) problem, which determines evolution by minimizing acceleration costs. TracingFlow directly regresses the acceleration field and initial velocity, avoiding ODE simulation. Our experiments show that TracingFlow achieves competitive and often improved reconstruction accuracy relative to existing simulation-free frameworks [42, 58, 59, 69], while effectively incorporating biological priors. Our contributions include: • We propose TracingFlow, a simulation-free TI framework based on second-order dynamics. It solves the proposed DOAT problem by regressing acceleration fields, reducing computational costs compared to ODE-based methods. • We provide theoretical guarantees by decoupling the design of the optimal transport plan and single-particle control. We further introduce an iterative strategy to transform the real-world Velocity Missed-DOAT (VM-DOAT) problem into a standard DOAT problem, enabling approximate solutions. • We demonstrate TracingFlow’s effectiveness on multi-time-point real world datasets. Results indicate superior distribution reconstruction accuracy and enhanced capability in preserving biological priors compared to state-of-the-art methods.

2

Related Works

Solving Optimal Transport via Flow Matching. Flow Matching [1, 35, 36, 48] transports probability distributions via learned flow maps, offering scalable simulation-free training. As an optimal control problem, Optimal Transport (OT) [47] decomposes into Optimal Coupling and single-particle geodesics, facilitating solutions via Flow Matching [15, 27, 50, 58]. This approach 2

extends to various variants and generalizations [4, 8, 13, 15, 26, 45, 46, 59, 61, 68]. However, these typically employ first-order dynamics. While second-order approaches have emerged [56], they require iterative training of the acceleration field to achieve optimal coupling. Moreover, they rely on heuristic methods for initial velocity estimation, rather than integrating it as a component of the optimal control problem. TracingFlow addresses this by decomposing the Dynamical Acceleration Optimal Transport (DOAT) problem into optimal coupling and single particle optimal trajectory, achieving a truly simulation-free process for DOAT. Overcoming Single-Valued Velocity Constraint in Flow Matching. In standard Flow Matching, the velocity field is modeled as a function solely dependent on x and t. This singlevalued dependency prevents the model from representing intersecting trajectories with distinct velocities, often resulting in curved inference paths and increased computational costs. To decouple the velocity from the current position, the velocity network can incorporate additional inputs like class labels [54, 73], initial positions [14], or hidden states [20]. Additionally, [69] employs hierarchical generation to mitigate this issue. However, these methods do not extend to second-order dynamics. While [42] utilized a second-order system, it is limited to constantacceleration paths. In contrast, TracingFlow offers a flexible second-order setting. By designing minimal-cost paths, it effectively resolves the limitation of the single-valued velocity field. Single-cell Trajectory Inference and Lineage Integration. Optimal Transport (OT) is a robust framework for inferring cellular dynamics from scRNA-seq data [2, 6, 10–12, 24, 25, 28, 29, 33, 37, 40, 41, 44, 51–53, 55, 57, 60, 65, 66, 70, 74]. However, transcriptomic similarity alone often fails to resolve complex trajectories. To mitigate this, studies incorporate lineage tracing priors such as clonal barcode information. Existing approaches integrate such info via regularization [17, 49], structural alignment [32], velocity mapping [62], or sparse transition modeling [18, 21, 63]. While improving accuracy, these generally focus on discrete couplings or static matrices. TracingFlow advances this by embedding priors into a continuous OT-based second-order flow matching framework, capturing complex dynamics consistent with lineage information.

3

Mathematical Background

In this section, we provide the mathematical formulation of the problem addressed by TracingFlow. Dynamical Optimal Acceleration Transport (DOAT) Problem Let X ⊂ Rd and V ⊂ Rd denote the position and velocity spaces, respectively. We define the augmented space as S := X × V, so that the full system state is s = (x, v) ∈ S. Consider the second-order dynamic system: ẋ = v v̇ = a(x, v, t), t ∈ [t0 , tK ]. (1) Modeling a large ensemble of such particles via a probability density ρt in S, we follow [5, 11] to formulate the following optimal control problem: min JDOAT (a, ρ) = a,ρ

t0

s.t.

1 ∥a(x, v, t)∥2 ρt (x, v) dxdvdt 2 S

Z tK Z

(2)

∂t ρt + ∇x · (vρt ) + ∇v · (aρt ) = 0,

(3)

ρti (x, v) = µi (x, v),

(4)

i = 0, . . . , K.

Here, constraints are imposed at K + 1 timestamps ti with given densities µi . Equation (3) represents the continuity equation in the augmented space. Analogous to standard Dynamical 3

Optimal Transport [47], we term this the Dynamical Optimal Acceleration Transport (DOAT) problem. To facilitate the subsequent numerical solution via particle methods, we assume the existence of an optimal flow map. Velocity-Missed Dynamical Optimal Acceleration Transport (VM-DOAT) Problem However, in practical applications such as single-cell datasets, only particle positions are observable, rendering velocities as unknown quantities. This presents a discrepancy with the DOAT formulation. Consequently, TracingFlow addresses a relaxation of the problem: min JVM-DOAT (a, ρ) = a,ρ

1 ∥a(x, v, t)∥2 ρt (x, v) dxdvdt S 2

Z tK Z t0

s.t.

(5)

∂t ρt (x, v) + ∇x · (vρt ) + ∇v · (aρt ) = 0, Z

(pos)

ρti (x, v) dv = µi

(6)

(x), i = 0, 1, . . . , K.

(7)

Here, the constraints are imposed only on the marginal distribution of positions, ρtj (x, v) dv, at each time point, with no constraints on the velocity distribution. We refer to this problem as Velocity Missed-Dynamical Optimal Acceleration Transport (VM-DOAT). R

4

Methodology of TracingFlow

In this section, we first introduce the key properties of the DOAT problem that are essential for our algorithm design. Subsequently, drawing inspiration from Flow Matching approaches for standard Dynamical OT, we devise an algorithm to exactly solve the DOAT problem by combining Optimal Coupling with single-particle optimal trajectories. Finally, we introduce a global path-based velocity completion scheme that effectively bridges the gap between VM-DOAT and the standard DOAT framework. By inferring missing velocities through the lens of global trajectory consistency, our approach enables the application of optimal transport theory to solve the VM-DOAT problem in a principled manner.

4.1

Key Properties of DOAT Problem

Optimal Single Particle Trajectory The DOAT problem is an optimal control problem over distributions. We begin by considering the corresponding single-particle optimal control problem. Specifically, the optimal trajectory is a cubic interpolation determined by the boundary states. Proposition 4.1. Let γ : [0, T ] → Rd be a twice continuously differentiable curve satisfying the boundary conditions γ(0) = x0 , γ ′ (0) = v0 , γ(T ) = xT , and γ ′ (T ) = vT . The variational problem Z T 1 ′′ min ∥γ (t)∥22 dt (8) γ 0 2 admits a unique minimizer, which is a cubic polynomial of the form γ(t) = c3 t3 + c2 t2 + c1 t + c0 ,

(9)

where the coefficients are detailed in Section A.1.This result is got by Euler-Lagrange formula. See a detailed proof in Section A.1.

4

Optimal Coupling By utilizing the single-particle optimal trajectory, we can immediately obtain the minimum cost for transporting a particle from (x0 , v0 ) to (xT , vT ): Corollary 4.2. Substitute Equation (9) to Equation (8), we get the minimum cost : 2 C[0,T ] = 3 (∥v0 ∥2 + ⟨v0 , vT ⟩ + ∥vT ∥2 )T 2 + 3⟨v0 + vT , x0 − xT ⟩T + 3∥x0 − xT ∥2 T 



(10)

The minimum cost allows us to define an optimal transport problem between t = 0 and t = T: Z ∗ π[0,T ] = min C[0,T ] [(x0 , v0 ), (xT , vT )]dπ[0,T ] (11) π[0,T ]

where π [(x , v ), (xT , vT )] is any coupling satisfying the constraints π[0,T ] dx0 dv0 = µT (x, v) R [0,T ] 0 0 and π[0,T ] dxT dvT = µ0 (x, v) , µ0 (x, v) and µT (x, v) are two distributions in augmented space. We term this formulation the Static Optimal Acceleration Transport (SOAT) problem. It is worth noting that while DOAT permits crossings in the projected position space, it prohibits trajectory crossing in the augmented space, which guarentees the neatness of trajectories. The following remark provides the mathematical justification for this non-crossing property in the augmented space, demonstrating that any such intersection implies a strictly higher transport cost. R

Remark 4.3. Consider two pairs of points in R2d , (x0 , v0 ), (xT , vT ) and (x̃0 , ṽ0 ), (x̃T , ṽT ) with no points overlapping. Denote the cubic interpolation curves between each pair by γ(t), γ̃(t)(t ∈ [0, T ]), respectively. If the curves cross with each other at time point t0 ∈ (0, T ), then we have C[0,T ] (x0 , v0 , xT , vT ) + C[0,T ] (x̃0 , ṽ0 , x̃T , ṽT ) > C[0,T ] (x0 , v0 , x̃T , ṽT ) + C[0,T ] (x̃0 , ṽ0 , xT , vT ). See a detailed proof in Section A.2. Equation (2) represents a multi-marginal optimal transport problem over the interval [t0 , tK ]. To address this, we formulate sequential SOAT problems between adjacent time steps ti and ti+1 by applying the substitutions x0 → xi , xT → xi+1 , v0 → vi , vT → vi+1 , and T → ti+1 − ti ∗ to Equation (11). For simplicity, let Ci→i+1 denote the resulting pairwise cost and πi→i+1 the optimal plan. Due to the Markovian nature of the dynamics, the full multi-marginal plan ∗ decomposes into the sequence of these local plans {πi→i+1 } for i = 0, . . . , K − 1.

4.2

Conditional and Marginal Probability Path

Obtaining Marginal Probability Path via Mixing Conditional Probability Paths Similar to the approach in standard Flow Matching [35, 58], directly constructing the Marginal Probability Path is difficult. Therefore, we condition on z and define the Marginal Probability Path as a mixture of Conditional Probability Paths: ρt (x, v) =

Z

ρt (x, v|z)q(z)dz

(12)

The condition z consists of K +1 points selected respectively from the snapshot data at K +1 time points t = 0, 1, · · · , K. If ρt (x, v|z) is generated by a conditional acceleration field a(x, v, t|z) starting from the initial condition ρ0 (x, v|z), then the marginal acceleration field: a(x, v, t) := Eq(z)



a(x, v, t|z)ρt (x, v|z) ρt (x, v)



(13)

will generate the marginal probability path ρt (x, v). Theorem 4.4. The marginal acceleration field in Eq. (13) generates the marginal probability path. See the proof in Section A.3. 5

Equivalence between Regressing Conditional Acceleration Fields and Marginal Acceleration Fields Assuming the marginal acceleration field a(x, v, t) is known and we can sample from the marginal probability path ρt (x, v), we can directly regress a(x, v, t) using a neural network. Let aθ (·, ·, ·) : Rd × Rd × R1 → Rd be a time-dependent acceleration field parameterized by a neural network with parameters θ. The acceleration flow matching loss is n

LAFM = E(x,v)∼ρt (x,v),t∼U [t0 ,tK ] ∥aθ − a(x, v, t)∥2

o

However, a(x, v, t) is generally intractable because it is defined via an expectation (integral), and its denominator ρt (x, v) also requires integration to compute. In contrast, the conditional acceleration field a(x, v, t|z) and the conditional probability path ρt (x, v|z) have simple forms. Therefore, in actual training, we use the following Conditional Acceleration Flow Matching (CAFM) objective: h

LCAFM = E (x,v)∼ρt (x,v|z), ∥aθ (x, v, t) − a(x, v, t|z)∥2

i

t∼U [t0 ,tK ],z∼q(z)

These two objectives are equivalent in training, as described by the following theorem: Theorem 4.5. If ρt (x, v) > 0 for any x ∈ Rd , v ∈ Rd , t ∈ [t0 , tK ], then LAFM and LCAFM differ only by a constant independent of θ. In other words, ∇θ LAFM = ∇θ LCAFM . The proof is left to Section A.4. Design of the Conditional Probability Path Consider the joint distributions µ0 (x, v), µ1 (x, v), ..., µK (x, v) at K + 1 time points. The optimal transport plans calculated via ∗ ∗ ∗ solving SOAT problem are π0→1 , π1→2 , · · · , πK−1→K . Taking the condition variable z = [(x0 , v0 ), · · · , (xK , vK )], and similar to standard Flow Matching, we choose a path with timevarying mean and invariant variance as the conditional probability path, i.e.: ρt (x, v|z) = N (x|mx (t), σx2 I) · N (v|mv (t), σv2 I)

(14)

Here, mx , mv are given by the Minimum Cost Trajectory described in Proposition 4.1, where for t ∈ [ti , ti+1 ]: mx (t) = c3 (t − ti )3 + c2 (t − ti )2 + c1 (t − ti ) + c0

mv (t)3c3 (t − ti )2 + 2c2 (t − ti ) + c1

, the coefficients are identical to those given in Section A.1, with x0 , v0 replaced by xi , vi , xT , vT replaced by xi+1 , vi+1 , and T replaced by ti+1 − ti . Based on the optimal transport plans, we define the optimal propagator (transition kernel) from time ti to ti+1 as the probability of a particle at (xi , vi ) at time ti being transported to (xi+1 , vi+1 ) at time ti+1 : ∗ Ki→i+1 [(xi , vi ), (xi+1 , vi+1 )] =

∗ πi→i+1 [(xi , vi ), (xi+1 , vi+1 )] µi (xi , vi )

(15)

Then, z ∼ µ0 (x0 , v0 ) ·

K−1 Y

∗ Ki→i+1 [(xi , vi ), (xi+1 , vi+1 )]

i=0

In the following, we demonstrate that the constructed marginal probability path recovers the true distribution at every given time point. Theorem 4.6. When σx → 0, σv → 0, the marginal probability density constructed by Equation (12) satisfies ρti (x, v) → µi (x, v) at the K + 1 given time points. Therefore, the Marginal Probability Path constructed by mixing Conditional Probability Paths can correctly reconstruct the marginal distributions. See the proof in Section A.5. Furthermore, we demonstrate that the design yields an exact solution to the DOAT problem. Proposition 4.7. As σx → 0 and σv → 0, the marginal probability path and the marginal acceleration field will converge to the solution to the DOAT problem. See the proof in Section A.6. 6

4.3

Transform VM-DOAT to DOAT

Relations between DOAT and VM-DOAT The optimal transport cost of the DOAT problem is a functional of the sequence of joint distributions µi (x, v) for i = 0, . . . , K. Formally, we denote this dependency as JDOAT [{µi (x, v)}|K i=0 ]. In the VM-DOAT setting, we only have R (pos) access to the marginal spatial distributions µi (x) = µi (x, v)dv. Consequently, the cost (pos) of VM-DOAT, denoted as JVM-DOAT [{µi (x)}|K i=0 ], relates to JDOAT through the following minimization: (pos)

JVM-DOAT [{µi

(pos)

(x)}] = min JDOAT [{µi ωi (v|x)

(x) · ωi (v|x)}],

(16)

where ωi (v|x) represents the conditional distribution of velocity, such that the joint distribution (pos) factorizes as µi (x, v) = µi (x) · ωi (v|x). Thus, our objective is to determine the optimal ωi (v|x) that minimizes the expression above, thereby transforming the VM-DOAT problem into a solvable DOAT instance. This transformation rests on two key properties: 1) The SOAT problem can be solved efficiently using the Sinkhorn method. Once the Transport Plan is computed, we can sample a trajectory [(x0 , v0 ), · · · , (xK , vK )]; 2) Given the positions x0:K of a trajectory, consider the following multi-marginal optimal control problem: min

Z tK t0

1 ′′ ∥γ (t)∥2 dt, 2

s.t. γ(ti ) = xi , i = 0, . . . , K.

(17)

The velocities of the optimal trajectory, denoted as V = [γ ′ (t0 ), · · · , γ ′ (tK )], can be obtained by solving a sparse linear system, as illustrated in the following proposition: Proposition 4.8. The velocities V defined above are the solution to the sparse linear system V A = B, where A ∈ R(K+1)×(K+1) is a tridiagonal matrix and B ∈ Rd×(K+1) . A detailed derivation and the explicit expressions for A and B are provided in Section A.7. It is worth noting that solving this linear system requires O(K) time, ensuring high computational efficiency. Iterative Scheme to Perform the Transformation Based on this decomposition, we employ an iterative scheme to estimate the conditional velocity distributions across t0 , . . . , tK . (0) We initialize ωi (v|x) = δ(v − 0) and subsequently alternate between the following two steps: • Solve SOAT Plans: Compute the K optimal SOAT plans between the distributions (pos) (m) ∗(m) µi (x)·ωi (v|x) for adjacent time points. This yields the plans πi→i+1 [(xi , vi ), (xi+1 , vi+1 )] for i = 0, . . . , K − 1, where m denotes the iteration index. This step effectively provides (m) the multi-marginal optimal transport plan given the current ωi (v|x). • Update Velocity Distributions: Update the conditional velocity distributions based on the computed SOAT plans: (m+1) ωi (v|xi ) =

R (m) q (z)δ(v̂i [x0 · · · xK ] − v) dz\i R ,

q (m) (z) dz\i

(18)

where dz\i denotes the volume element dx0 dv0 . . . dxK dvK excluding dxi (i.e., integrating over all variables except xi ). The term v̂i [x0 , . . . , xK ] denotes the optimal velocity at time (m+1) ti derived from the sequence x0 , . . . , xK via Proposition 4.8. Intuitively, ωi (v|xi ) is set to the distribution induced by the velocities of all optimal control trajectories passing through xi at time ti . Since the total transport cost is the sum of costs along individual trajectories, aligning v with the optimal velocity of each trajectory minimizes the global objective. 7

In the update step above, the trajectory variable z = [(x0 , v0 ), . . . , (xK , vK )] follows the joint distribution q (m) (z) defined as: (pos)

q (m) (z) = µ0

(m)

(x0 ) · ω0 (v0 |x0 ) ×

K−1 Y

∗(m)

Ki→i+1 [(xi , vi ), (xi+1 , vi+1 )],

i=0 ∗(m)

where the propagator Ki→i+1 is given by: ∗(m)

∗(m)

Ki→i+1 [(xi , vi ), (xi+1 , vi+1 )] =

πi→i+1 [(xi , vi ), (xi+1 , vi+1 )] (pos)

µi

(m)

(xi ) · ωi

(vi |xi )

(pos)

(m)

.

During the iteration process, the objective LDOAT [µi (x) · ωi (v|x)] decreases monotonically. Since the cost is bounded below by 0, the iterative algorithm is guaranteed to converge.

5

TracingFlow Algorithm for Trajectory Inference

In the practical implementation of TracingFlow, we first perform an approximation of the conditional velocity distribution to reduce the VM-DOAT problem to a DOAT problem. Subsequently, we employ flow matching to solve for the acceleration field of the DOAT problem. Notably, for lineage tracing data, we incorporate biological priors into this framework.

5.1

Approximation of Conditional Velocity Distribution

Determining the full conditional velocity distribution ωi (v|x) to transform VM-DOAT to DOAT (Section 4.3) is computationally prohibitive, potentially requiring an auxiliary generative model. We therefore adopt a deterministic approximation, assigning a single velocity to each data point. Consequently, Equation (18) is modified to: (m+1) (m+1) ωi (v|xi ) = δ(v − vexp (xi )), (m+1)

Intuitively, this assigns vexp through xi at time ti .

5.2

(m+1) vexp (xi ) =

R (m) q (z)v̂i [x0 · · · xK ] dz\i R .

q (m) (z) dz\i

(19)

(xi ) as the mean velocity of all optimal control trajectories passing

MiniBatch-OT

Solving SOAT repeatedly is computationally intensive for large-scale datasets. However, since the conditional velocity distribution ωi (v|x) is approximated as a Dirac delta, the joint empirical distribution µi (x, v) effectively becomes a superposition of Dirac deltas. This allows us to employ (n) a minibatch strategy: we partition µi and µi+1 into B corresponding minibatches {µi }B n=1 (n) ∗(n) and {µi+1 }B . By solving the local plans π independently and aggregating them as n=1 i→i+1 ∗(n)

∗ πi→i+1 = ⊕B n=1 πi→i+1 , we significantly enhance computational efficiency.

5.3

Incorporation of Biological Priors

In lineage tracing applications, let indices li and li+1 denote individual samples at time points ti and ti+1 , respectively. Each data point (xi,li , vi,li ) is associated with a barcode bi,li ∈ N+ . We incorporate this prior by modifying the cost matrix Ci→i+1 . Specifically, when computing the cost between the l-th sample at ti and the k-th sample at ti+1 , we compare their barcodes: if bi,li = ̸ bi+1,li+1 , the standard cost is multiplied by a penalty factor p0 (p0 > 1). This soft constraint discourages transitions between distinct lineages, effectively embedding biological knowledge into the dynamics learning process. 8

Table 1: Average W1 and W2 distances between the generated and ground truth distributions at various time points on the 2D Simulation , Cite 5D, and Cite 100D datasets for different algorithms. Method OT-CFM SF2 M 3MSBM MMFM HRF CAF TF (Ours)

5.4

Dynamics

2D Simulation

Cite 5D

Cite 100D

W1

W2

W1

W2

W1

W2

1st Order

0.5270 ± 0.0612 0.9926 ± 0.1857

0.5922 ± 0.0813 1.1318 ± 0.1734

0.8298 ± 0.0162 0.7650 ± 0.0375

0.9155 ± 0.0160 0.8956 ± 0.0408

10.5498 ± 0.0291 12.3426 ± 0.0735

10.6220 ± 0.0310 12.4757 ± 0.0783

2nd Order

1.1855 ± 0.2636 2.3689 ± 0.1830 1.5390 ± 0.1130 1.6385 ± 0.2777 0.3748 ± 0.0505

1.2767 ± 0.2215 2.6466 ± 0.1769 1.6631 ± 0.1214 1.6431 ± 0.2766 0.5180 ± 0.0360

3.6369 ± 0.3508 2.5346 ± 0.1520 1.4016 ± 0.0437 4.5589 ± 0.1073 0.5861 ± 0.0688

3.7703 ± 0.3178 2.7576 ± 0.1877 1.5313 ± 0.0430 4.7411 ± 0.1620 0.6386 ± 0.0775

15.7708 ± 0.1473 12.5353 ± 0.6268 10.1532 ± 0.0643 16.7429 ± 0.4371 8.0700 ± 0.1537

15.8833 ± 0.1832 12.7868 ± 0.7061 10.2336 ± 0.0707 16.9837 ± 0.4324 8.2530 ± 0.1889

Algorithm of TracingFlow

The TracingFlow algorithm comprises three main steps: First, we apply the iterative approximation from Section 5.1 to determine initial velocities, effectively transforming the VM-DOAT problem into a DOAT problem. Second, we compute the Optimal Transport Plan for SOAT as defined in Equation (11). Third, we sample conditional probability paths to train the acceleration network aθ (x, v, t) via regression. Additionally, to provide initial conditions for inference, we parameterize the unique initial velocities derived in the first step using a neural network v0,ξ (x), trained similarly via regression. The algorithm’s pseudocode is provided in Section D. Theoretically, TracingFlow recovers the exact distribution if the acceleration and initial velocity are perfectly fitted. In practice, where losses are non-zero, we prove that the positional W2 distance between the generated and true distributions is bounded by the sum of the velocity and AFM losses, and the relevant theorem is detailed in Section A.8.

6

Experiment Results

To evaluate the effectiveness of our algorithm, TracingFlow, we conducted three categories of experiments: 1) verifying whether the acceleration field of TracingFlow can transport the initial R (pos) (pos) data distribution µ0 (x) = ρt0 (x, v)dv to the data distribution µi (x) at other given times ti ; 2) verifying whether TracingFlow can effectively perform distribution interpolation and extrapolation; and 3) verifying whether TracingFlow can more effectively infer the dynamical laws in Lineage Tracing data given biological prior knowledge. Distribution Transport We evaluated TracingFlow on a 2-dimensional synthetic dataset and real single-cell omics datasets of varying dimensions. The evaluation metrics were the 1-Wasserstein (W1 ) and 2-Wasserstein (W2 ) distances, measuring the discrepancy between the fitted and ground truth distributions. The second-order dynamics model of TracingFlow provides stronger expressive power than first-order algorithms, yielding lower average W1 and W2 values across all datasets (Table 1). We visualized the evolutionary trajectories learned by OT-CFM and TracingFlow on the synthetic dataset in Figure 2. TracingFlow’s second-order model allows it to learn trajectories that cross in the position space x1 , x2 , whereas OT-CFM is unable to do so. In Section C.1, we present the distribution reconstruction accuracy at each time point, as well as experimental results on additional datasets (such as the Mexico Gulf dataset). Interpolation and Extrapolation To verify TracingFlow’s interpolation and extrapolation capabilities, we conducted a Hold-One-Out experiment on the 5-dimensional EB dataset (time points t = 0, . . . , 4). In each iteration, we withheld one time point from the latter four for training and calculated the W1 and W2 distances between the interpolated and true distributions. As

9

a) Traj. Generated by OT-CFM

b) Traj. Generated by Tracing Flow

Figure 2: On the Simulation 2D dataset: a) non-crossing paths learned by OT-CFM, and b) crossing paths learned by TracingFlow. Table 2: Average W1 and W2 distances between the predicted and ground truth distributions at heldout time points on the EB 5D dataset for different algorithms.

Table 3: Average lineage-weighted W1 and W2 distances between the generated and ground truth distributions at various time points on the 3D Simulation Lineage and Hematopoiesis datasets for different algorithms. Method

EB 5D

Method

Dynamics W1

W2

OT-CFM SF2 M

1st Order

5.1205 ± 0.2400 4.5602 ± 0.1041

5.6910 ± 0.2561 4.9905 ± 0.1262

3MSBM MMFM TF (Ours)

2nd Order

7.1901 ± 0.8667 9.7777 ± 0.9006 4.4974 ± 0.0950

7.4697 ± 0.8588 10.5980 ± 0.9407 4.7859 ± 0.1857

Dynamics

OT-CFM SF2 M 3MSBM MMFM HRF CAF TF w/o Bio-Prior TF (Ours)

3D Sim-Lineage

Hematopoiesis

W1

W2

W1

W2

1st Order

2.2383 ± 0.0205 1.5255 ± 0.0513

2.3162 ± 0.0216 1.5435 ± 0.0515

14.9945 ± 0.3442 18.7170 ± 0.1124

15.3971 ± 0.3180 18.9756 ± 0.1168

2nd Order

2.3126 ± 0.4756 2.6098 ± 0.3149 1.7934 ± 0.4043 2.6806 ± 0.2777 1.5292 ± 0.1714 0.4538 ± 0.2039

2.4905 ± 0.4582 3.0309 ± 0.3265 1.8085 ± 0.4022 2.7547 ± 0.3521 1.5532 ± 0.2303 0.5328 ± 0.2066

18.3313 ± 0.9834 27.0830 ± 2.9299 15.4571 ± 0.1266 15.2165 ± 0.0685 16.6166 ± 0.1229 14.0470 ± 0.1530

18.6427 ± 0.9413 29.5421 ± 3.1420 15.7138 ± 0.1282 15.4974 ± 0.0576 17.0110 ± 0.1456 14.4589 ± 0.1815

shown in Table 2, TracingFlow achieved the highest accuracy, demonstrating that models based on second-order dynamics are superior in capturing high-curvature and non-linear trajectories. Lineage Tracing Data To verify whether TracingFlow better handles lineage tracing data given biological priors, we experimented on a 3D Simulation Lineage Dataset and the Hematopoiesis Dataset. Following Section 5.3, we introduced biological priors into TracingFlow by modifying the transport cost matrix, other baselines also incorporates priors via conditional velocity fields. We used lineage-weighted W1 and W2 distances to evaluate consistency with biological priors while learning dynamics. Results in Table 3 show that incorporating priors enables TracingFlow to outperform other algorithms. In Figure 3, we visualized dynamics on the simulation dataset to demonstrate how TracingFlow learns trajectories conforming to priors. In Figure 4, we plotted the PCA-reduced positions of two barcodes from the Hematopoiesis Dataset at t = 1; the positions inferred by TracingFlow are significantly closer to the ground truth than those by OT-CFM.

7

Conclusion and Limitation

In this work, we introduce TracingFlow, a Flow Matching framework for the Dynamical Acceleration Optimal Transport (DOAT) problem. By pre-solving SOAT to learn acceleration and initial velocity, it enables efficient, simulation-free solutions for large-scale DOAT problem. Experiments on simulated and real world data demonstrate that second-order dynamics enhance expressivity, yielding precise temporal distribution recovery. Additionally, incorporating biological priors better preserves intercellular lineages. A limitation is the computationally intensive SOAT pre-calculation, though this is mitigable via minibatch-OT. Furthermore, explicitly modeling the theoretical conditional velocity distribution per data point is prohibitive, so we approximate it using its expectation. The objective 10

a)Traj. Generated by OT-CFM

b) Traj. Generated by Tracing Flow

a) Data points Generated by OT-CFM

(t=1)

Wrong Barcodes

b) Data points Generated by Tracing Flow

(t=1)

Correct Barcodes

Figure 3: On the Sim-Lineage dataset: a) paths learned by OT-CFM without considering biological priors, and b) paths learned by TracingFlow that preserve biological priors.

Figure 4: Comparison of generated (circles) and ground truth (’x’) positions at t = 1 for two barcodes on the hematopoiesis dataset [64]. (a) OT-CFM. (b) TracingFlow. Visualized via 2D PCA.

RT 1

2 0 2 ∥a∥ dt also lacks a clear physical interpretation, as standard classical actions exclude second

time derivatives (see Section E.2). Despite current sequencing measurements lacking velocity data, priors suggest biological systems may follow higher-order dynamics due to the existence of complex regulations. Future work will extend TracingFlow to further model these dynamics.

References [1] Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. arXiv preprint arXiv:2209.15571, 2022. [2] Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797, 2023. [3] Vladimir Igorevich Arnol’d. Mathematical methods of classical mechanics, volume 60. Springer Science & Business Media, 2013. [4] 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. [5] Jean-David Benamou, Thomas O Gallouët, and François-Xavier Vialard. Second-order models for optimal transport and cubic splines on the wasserstein space. Foundations of Computational Mathematics, 19(5):1113–1143, 2019. [6] Charlotte Bunne, Ya-Ping Hsieh, Marco Cuturi, and Andreas Krause. The schrödinger bridge between gaussian measures has a closed form. In International Conference on Artificial Intelligence and Statistics, pages 5802–5833. PMLR, 2023. [7] 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. [8] Zihan Cao, Yu Zhong, and Liang-Jian Deng. Taming flow matching with unbalanced optimal transport into fast pansharpening. arXiv preprint arXiv:2503.14975, 2025. [9] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018. 11

[10] Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou. Likelihood training of schrödinger bridge using forward-backward SDEs theory. In International Conference on Learning Representations, 2022. [11] Yongxin Chen, Giovanni Conforti, Tryphon T Georgiou, and Luigia Ripani. Multi-marginal schrödinger bridges. In International Conference on Geometric Science of Information, pages 725–732. Springer, 2019. [12] 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. [13] 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. [14] Valentin De Bortoli, Guan-Horng Liu, Tianrong Chen, Evangelos A Theodorou, and Weilie Nie. Augmented bridge matching. arXiv preprint arXiv:2311.06978, 2023. [15] Luca Eyring, Dominik Klein, Théo Uscidda, Giovanni Palla, Niki Kilbertus, Zeynep Akata, and Fabian J Theis. Unbalancedness in neural Monge maps improves unpaired domain translation. In The Twelfth International Conference on Learning Representations, 2024. [16] Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, et al. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1–8, 2021. [17] Aden Forrow and Geoffrey Schiebinger. Lineageot is a unified framework for lineage tracing and trajectory inference. Nature Communications, 12(1), 2021. [18] Mingze Gao, Melania Barile, Shirom Chabra, Myriam Haltalli, Emily F. Calderbank, Yiming Chao, Weizhong Zheng, Nicola K. Wilson, Elisa Laurenti, Berthold Göttgens, and Yuanhua Huang. CLADES: a hybrid neuralode-gillespie approach for unveiling clonal cell fate and differentiation dynamics. Nature Communications, 16(1), 2025. [19] Gennady Gorin, Valentine Svensson, and Lior Pachter. Protein velocity and acceleration from single-cell multiomics experiments. Genome biology, 21(1):39, 2020. [20] Pengsheng Guo and Alexander G Schwing. Variational rectified flow matching. arXiv preprint arXiv:2502.09616, 2025. [21] Wenbo Guo, Zeyu Chen, Xinqi Li, Jingmin Huang, Qifan Hu, and Jin Gu. sctrace+: Enhancing cell fate inference by integrating the lineage-tracing and multi-faceted transcriptomic similarity information. Cell Systems, 16(9):101398, 2025. [22] Peter Halmos, Xinhao Liu, Julian Gold, Feng Chen, Li Ding, and Benjamin J Raphael. Dest-ot: Alignment of spatiotemporal transcriptomics data. Cell Systems, 16(2), 2025. [23] Ari Hong, Sangseon Lee, and Kwangsoo Kim. Multi-omic relay velocity modeling uncovers dynamic chromatin-transcription regulation across cell states. Nature Communications, 2025. [24] Guillaume Huguet, Daniel Sumner Magruder, Alexander Tong, Oluwadamilola Fasina, Manik Kuchroo, Guy Wolf, and Smita Krishnaswamy. Manifold interpolating optimaltransport flows for trajectory inference. Advances in neural information processing systems, 35:29705–29718, 2022. 12

[25] Qi Jiang and Lin Wan. A physics-informed neural SDE network for learning cellular dynamics from time-series scRNA-seq data. Bioinformatics, 40:ii120–ii127, 09 2024. ISSN 1367-4811. [26] 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. [27] Dominik Klein, Théo Uscidda, Fabian Theis, and Marco Cuturi. Genot: Entropic (gromov) wasserstein flow matching with applications to single-cell genomics. Advances in Neural Information Processing Systems, 37:103897–103944, 2024. [28] Dominik Klein, Giovanni Palla, Marius Lange, Michal Klein, Zoe Piran, Manuel Gander, Laetitia Meng-Papaxanthos, Michael Sterr, Lama Saber, Changying Jing, et al. Mapping cells through time and space with moscot. Nature, pages 1–11, 2025. [29] Takeshi Koshizuka and Issei Sato. Neural lagrangian schrödinger bridge: Diffusion modeling for population dynamics. In The Eleventh International Conference on Learning Representations, 2023. [30] Kai Kretzschmar and Fiona M. Watt. Lineage tracing. Cell, 148(1–2):33–45, 2012. [31] 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, pages 2022–04, 2022. [32] Marius Lange, Zoe Piran, Michal Klein, Bastiaan Spanjaard, Dominik Klein, Jan Philipp Junker, Fabian J. Theis, and Mor Nitzan. Mapping lineage-traced cells across time points with moslin. Genome Biology, 25(1), 2024. [33] Hugo Lavenant, Stephen Zhang, Young-Heon Kim, Geoffrey Schiebinger, et al. Toward a mathematical theory of trajectory inference. The Annals of Applied Probability, 34(1A): 428–500, 2024. [34] Chen Li, Maria C Virgilio, Kathleen L Collins, and Joshua D Welch. Multi-omic single-cell velocity models epigenome–transcriptome interactions and improves cell fate prediction. Nature biotechnology, 41(3):387–398, 2023. [35] 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 Representations, 2023. [36] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003, 2022. [37] Suryanarayana Maddu, Victor Chardès, Michael Shelley, et al. Inferring biological processes with intrinsic noise from cross-sectional data. arXiv preprint arXiv:2410.07501, 2024. [38] Shanjun Mao, Chenyang Zhang, Runjiu Chen, Shan Tang, Xiaodan Fan, and Jie Hu. Cell lineage tracing: Methods, applications, and challenges. Quantitative Biology, 13(4), 2025. [39] Kevin R Moon, David Van Dijk, Zheng Wang, Scott Gigante, Daniel B Burkhardt, William S Chen, Kristina Yim, Antonia van den Elzen, Matthew J Hirn, Ronald R Coifman, et al. Visualizing structure and transitions in high-dimensional biological data. Nature biotechnology, 37(12):1482–1492, 2019. 13

[40] Kirill Neklyudov, Rob Brekelmans, Daniel Severo, and Alireza Makhzani. Action matching: Learning stochastic dynamics from samples. In International conference on machine learning, pages 25858–25889. PMLR, 2023. [41] Kirill Neklyudov, Rob Brekelmans, Alexander Tong, Lazar Atanackovic, Qiang Liu, and Alireza Makhzani. A computational framework for solving wasserstein lagrangian flows. In Forty-first International Conference on Machine Learning, 2024. [42] Dogyun Park, Sojin Lee, Sihyeon Kim, Taehoon Lee, Youngjoon Hong, and Hyunwoo J Kim. Constant acceleration flow. Advances in Neural Information Processing Systems, 37: 90030–90060, 2024. [43] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017. [44] Qiangwei Peng, Peijie Zhou, and Tiejun Li. stVCR: Spatiotemporal dynamics of single cells. bioRxiv, pages 2024–06, 2024. [45] Qiangwei Peng, Zihan Wang, Junda Ying, Yuhao Sun, Qing Nie, Lei Zhang, Tiejun Li, and Peijie Zhou. WFR-FM: Simulation-free dynamic unbalanced optimal transport. arXiv preprint arXiv:2601.06810, 2026. [46] Katarina Petrović, Lazar Atanackovic, Kacper Kapusniak, Michael M. Bronstein, Joey Bose, and Alexander Tong. Curly flow matching for learning non-gradient field dynamics. In Learning Meaningful Representations of Life (LMRL) Workshop at ICLR 2025, 2025. [47] Gabriel Peyré and Marco Cuturi. Computational optimal transport: With applications to data science. Now Foundations and Trends, 2019. [48] Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky TQ Chen. Multisample flow matching: Straightening flows with minibatch couplings. arXiv preprint arXiv:2304.14772, 2023. [49] Neha Prasad, Karren Yang, and Caroline Uhler. Optimal transport using gans for lineage tracing. arXiv preprint arXiv:2007.12098, 2020. [50] 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. [51] Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 176(4):928–943, 2019. [52] 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. [53] Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching. Advances in Neural Information Processing Systems, 36, 2024. [54] Sagar Shrestha and Xiao Fu. Diversified flow matching with translation identifiability. arXiv preprint arXiv:2511.05558, 2025. 14

[55] Yuhao Sun, Zhenyi Zhang, Zihan Wang, Tiejun Li, and Peijie Zhou. Variational regularized unbalanced optimal transport: Single network, least action. arXiv preprint arXiv:2505.11823, 2025. [56] Panagiotis Theodoropoulos, Augustinos D Saravanos, Evangelos A Theodorou, and GuanHorng Liu. Momentum multi-marginal schr\" odinger bridge matching. arXiv preprint arXiv:2506.10168, 2025. [57] 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, pages 9526–9536. PMLR, 2020. [58] 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, 2024. ISSN 2835-8856. Expert Certification. [59] 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, pages 1279–1287. PMLR, 2024. [60] Elias Ventre, Aden Forrow, Nitya Gadhiwala, Parijat Chakraborty, Omer Angel, and Geoffrey Schiebinger. Trajectory inference for a branching sde model of cell differentiation. arXiv preprint arXiv:2307.07687, 2023. [61] 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. [62] Kun Wang, Liangzhen Hou, Xin Wang, Xiangwei Zhai, Zhaolian Lu, Zhike Zi, Weiwei Zhai, Xionglei He, Christina Curtis, Da Zhou, and Zheng Hu. Phylovelo enhances transcriptomic velocity field mapping using monotonically expressed genes. Nature Biotechnology, 42(5): 778–789, 2023. [63] Shou-Wen Wang, Michael J. Herriges, Kilian Hurley, Darrell N. Kotton, and Allon M. Klein. Cospar identifies early cell fate biases from single-cell transcriptomic and lineage information. Nature Biotechnology, 40(7):1066–1074, 2022. [64] Caleb Weinreb, Alejo Rodriguez-Fraticelli, Fernando D Camargo, and Allon M Klein. Lineage tracing on transcriptional landscapes links state to fate during differentiation. Science, 367(6479):eaaw3381, 2020. [65] Maosheng Yang. Topological schrödinger bridge matching. In The Thirteenth International Conference on Learning Representations, 2025. [66] Grace Hui Ting Yeo, Sachit D Saksena, and David K Gifford. Generative modeling of single-cell time series with prescient enables prediction of cell trajectories with interventions. Nature communications, 12(1):3222, 2021. [67] Jiaqi Zhang, Erica Larschan, Jeremy Bigness, and Ritambhara Singh. scNODE: generative model for temporal single cell transcriptomic data prediction. Bioinformatics, 40 (Supplement_2):ii146–ii154, 09 2024. ISSN 1367-4811. [68] Xi Nicole Zhang, Yuan Pu, Yuki Kawamura, Andrew Loza, Yoshua Bengio, Dennis Shung, and Alexander Tong. Trajectory flow matching with applications to clinical time series modelling. Advances in Neural Information Processing Systems, 37:107198–107224, 2024. 15

[69] Yichi Zhang, Yici Yan, Alex Schwing, and Zhizhen Zhao. Towards hierarchical rectified flow. arXiv preprint arXiv:2502.17436, 2025. [70] Zhenyi Zhang, Tiejun Li, and Peijie Zhou. Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport. In The Thirteenth International Conference on Learning Representations, 2025. [71] Zhenyi Zhang, Yuhao Sun, Qiangwei Peng, Tiejun Li, and Peijie Zhou. Integrating dynamical systems modeling with spatiotemporal scRNA-Seq data analysis. Entropy, 27(5), 2025. ISSN 1099-4300. [72] Grace XY Zheng, Jessica M Terry, Phillip Belgrader, Paul Ryvkin, Zachary W Bent, Ryan Wilson, Solongo B Ziraldo, Tobias D Wheeler, Geoff P McDermott, Junjie Zhu, et al. Massively parallel digital transcriptional profiling of single cells. Nature communications, 8 (1):14049, 2017. [73] Qunxi Zhu and Wei Lin. Switched flow matching: Eliminating singularities via switching odes. arXiv preprint arXiv:2405.11605, 2024. [74] Qunxi Zhu, Bolin Zhao, Jingdong Zhang, Peiyang Li, and Wei Lin. Governing equation discovery of a complex system from snapshots. arXiv preprint arXiv:2410.16694, 2024.

16

Contents of Appendix Proofs of Main Theorems . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 17

A

A.1

Proof of Proposition 4.1 : Optimal Trajectory of Single Particle . . . . . . . . . . . . . . . . .17

A.3

Proof of Remark 4.4 : No-Cross Property . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .18

A.4

Proof of Theorem 4.5 : Marginal Acceleration Field . . . . . . . . . . . . . . . . . . . . . . . . . . . . 19

A.5

Proof of Theorem 4.6 : The Relation Between AFM Loss and CAFM Loss . . . . . . 19

A.6

Proof of Theorem 4.7 : Marginal Probability Path . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 20

A.7

Proof of Proposition 4.8 : TracingFlow Exactly Solves DOAT Problem . . . . . . . . . . 21

A.8

Proof of Proposition 4.9 : Solve Linear System to Get Optimal Velocity . . . . . . . . 22

A.9

Proof of Theorem A.1 : W2 Distance can be bounded by LAFM and Lv0 . . . . . . . . 24

Datasets and Evaluation Metric . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 26

B

B.1

Experiment Setup . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 26

B.2

Evaluation Metrics . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 27

B.3

Datasets . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 27

Experiment Details . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 28

C

C.1

Experiment Details across Various Datasets . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 28

C.2

Training Time and Scalability of Tracing Flow . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 32

C.3

Sensitivity Analysis on Minibatch-OT . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 32

D

Pseudocode for Algorithm . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 34

E

Discussion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 35 E.1

Relation with other Algorithms . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 35

E.2

Is S = X × V a phase space? . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 35

A

Proofs of Main Theorems

A.1

Proof of Proposition 4.1

Proof. First, we show that the unique minimizer is a cubic polynomial. Consider the energy functional defined on the interval [0, T ]: J [γ] =

1 2

Z T 0

∥γ ′′ (t)∥22 dt.

(20)

The integrand L(γ, γ ′ , γ ′′ ) = 12 ∥γ ′′ ∥22 depends on derivatives up to the second order. The necessary condition for optimality is given by the Euler–Lagrange equation: ∂L d − ∂γ dt



∂L ∂γ ′



d2 + 2 dt



∂L ∂γ ′′



= 0.

(21)

Since ∂L/∂γ = 0, ∂L/∂γ ′ = 0, and ∂L/∂γ ′′ = γ ′′ , the equation simplifies to: d2 ′′ γ (t) = γ (4) (t) = 0. dt2

(22)

This implies that the optimal trajectory γ(t) is a cubic polynomial: γ(t) = c3 t3 + c2 t2 + c1 t + c0 . 17

(23)

Imposing the boundary conditions γ(0) = x0 , γ ′ (0) = v0 and γ(T ) = xT , γ ′ (T ) = vT leads to the linear system:   c0 = x0 ,     c = v , 1 0 (24) 3 + c T2 + c T + c = x ,  c T  3 2 1 0 T    3c T 2 + 2c T + c = v . 3 2 1 T Solving for the coefficients yields the unique solution:   c0 = x0 ,       c1 = v0 ,

3(xT − x0 ) − (2v0 + vT )T

(25)

c2 = ,    T2    (v + vT )T − 2(xT − x0 )  c3 = 0 . T3

A.2

Proof of Remark 4.3

Proof. First we propose a lemma: Lemma. Let C[0,t] (x0 , v0 , xt , vt ) be the least acceleration cost through [0, t] ,(x0 , v0 ), (xT , vT ) ∈ X × V are two different points. Then for arbitrary midpoint τ ∈ (0, 1) and state (x, v) ∈ X × V, the following inequality holds: C[0,τ ] (x0 , v0 , x, v) + C[τ,T ] (x, v, xT , vT ) ≥ C[0,T ] (x0 , v0 , xT , vT )

(26)

The "=" holds if and only if (x, v) = γ ∗ (τ ), where γ ∗ (t)(t ∈ [0, T ]) is the cubic interpolation trajectory connecting (x0 , v0 ) and (xT , vT ). Proof. Fix τ ∈ (0, T ). Define objective function Fτ (x, v) = C[0,τ ] (x0 , v0 , x, v)+C[τ,T ] (x, v, xT , vT ). By Corollary 4.2 , we can rewrite Fτ (x, v) =

o 2 n 2 2 2 2 (∥v ∥ + ⟨v , v⟩ + ∥v∥ )τ + 3⟨v + v, x − x⟩τ + 3∥x − x∥ + 0 0 0 0 0 τ3

n o 2 (∥vT ∥2 + ⟨vT , v⟩ + ∥v∥2 )(T − τ )2 + 3⟨vT + v, x − xT ⟩T + 3∥xT − x∥2 3 (T − τ )

(27)

Which is a quadratic form of (x, v) up to a constant. Fτ (x, v) ≥ 0 always holds, namely, the quadratic form has a finite lower bound, so it is semi-quadratic. To minimize Fτ (x, v), we can take differentiations: −(v0 + v)τ + 2(x − x0 ) (v + vT )(T − τ ) + 2(x − xT ) ∇x Fτ (x, v) = 3 + τ2 (T − τ )2 

∇v Fτ (x, v) =



= 0.

(v0 + 2v)τ + 3(x0 − x) (vT + 2v)(T − τ ) + 3(x − xT ) + = 0, τ2 (T − τ )2

(28) (29)

Denote the cubic interpolation from (x0 , v0 ) to (x, v) on [0, τ ] and (x, v) to (xT , vT ) on [τ, T ] by γ[0,τ ] (t) and γ[τ,T ] (t), respectively.The coefficients are known according to Section A.1. Then we have: ′′

γ[0,τ ] (τ − ) = 6·

(v0 + v)τ − 2(x − x0 ) 3(x − x0 ) − (2v0 + v)τ 3(x0 − x) + (v0 + 2v)τ +2· = 2· , 2 2 τ τ τ2 (30)

18

′′

γ[τ,T ] (τ + ) = 6 ·

(v + vT )τ − 2(xT − x) 3(xT − x) − (2v + vT )(T − τ ) +2· τ2 (T − τ )2 3(x − xT ) + (v + 2vT )(T − τ ) =2· (T − τ )2

(31)

(v0 + v)τ − 2(x − x0 ) τ2

(32)

(v + vT )(T − τ ) − 2(xT − x) (T − τ )2

(33)

′′′

γ[0,τ ] (τ − ) = 6 · ′′′

γ[τ,T ] (τ + ) = 6 ·

′′′

′′′

′′

′′

Then we can find that 28 and 29 appear to be γ[0,τ ] (τ − ) = γ[τ,T ] (τ + ), γ[0,τ ] (τ − ) = γ[τ,T ] (τ + ). (i)

(i)

Because the two curves are cubic, and γ[0,τ ] (τ − ) = γ[τ,T ] (τ + ), i = 0, 1, 2, 3, so we know that in the optimal case γ[0,τ ] = γ ∗ [0,τ ] , γ[τ,T ] = γ ∗ [τ,T ] .. That is, the unique solution of 28 and 29 is ′

(γ ∗ (τ ), γ ∗ (τ )), which completes the proof. Back to our target, assume γ and γ̃ cross at (xτ , vτ ), then we have C[0,T ] (x0 , v0 , xT , vT ) + C[0,T ] (x̃0 , ṽ0 , x̃T , ṽT ) n

o

n

o

n

o

n

o

= C[0,τ ] (x0 , v0 , xτ , vτ ) + C[τ,T ] (xτ , vτ , xT , vT ) + C[0,τ ] (x̃0 , ṽ0 , xτ , vτ ) + C[τ,T ] (xτ , vτ , x̃T , ṽT ) = C[0,τ ] (x0 , v0 , xτ , vτ ) + C[τ,T ] (xτ , vτ , x̃T , ṽT ) + C[0,τ ] (x̃0 , ṽ0 , xτ , vτ ) + C[τ,T ] (xτ , vτ , xT , vT ) ≥C[0,T ] (x0 , v0 , x̃T , ṽT ) + C[0,T ] (x̃0 , ṽ0 , xT , vT ).

(34) The equality holds if and only if γ and γ̃ coincide. Consequently, the strict inequality holds for distinct trajectories, which implies the non-crossing property. This completes the proof.

A.3

Proof of Theorem 4.4

Proof. Consider the time derivative of the marginal distribution: ∂t ρt (x, v) = ∂t =

Z

=−

Z

ρt (x, v|z)q(z)dz



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

∇x · (ρt (x, v|z)v)q(z)dz −

= −∇x ·

Z

Z

∇v · (ρt (x, v|z)a(x, v, t|z))q(z)dz



ρt (x, v|z)vq(z)dz − ∇v ·

Z

ρt (x, v|z)a(x, v, t|z)q(z)dz

= −∇x · (ρt (x, v)v) − ∇v · (at (x, v)ρt (x, v))



(35)

Therefore, ρt (x, v) satisfies the continuity equation: ∂t ρt (x, v) = −∇x · (ρt (x, v)v) − ∇v · (at (x, v)ρt (x, v))

(36)

This completes the proof.

A.4

Proof of Theorem 4.5

Proof. Consider the gradient of the AFM Loss and CAFM Loss: ∇θ LAFM = ∇θ Eρt (x,v) (∥aθ (x, v, t)∥2 − 2aTθ (x, v, t)a(x, v, t)) ∇θ LCAFM = ∇θ Eρt (x,v|z)q(z) (∥aθ (x, v, t)∥ 19

2

− 2aTθ (x, v, t)a(x, v, t|z))

(37) (38)

For the first term of ∇θ LAFM : Eρt (x,v) ∥aθ (x, v, t)∥2 =

Z

ρt (x, v)∥aθ (x, v, t)∥2 dxdv

=

Z

ρt (x, v|z)q(z)∥aθ (x, v, t)∥2 dxdvdz

= Eρt (x,v|z)q(z) ∥aθ (x, v, t)∥2

(39)

For the second term: Eρt (x,v) [aTθ (x, v, t)a(x, v, t)] =

Z

ρt (x, v)aTθ (x, v, t)a(x, v, t)dxdv

=

Z

ρt (x, v)aTθ (x, v, t)

=

Z Z

Z

a(x, v, t|z)ρt (x, v|z) q(z) dz dxdv ρt (x, v) 

aTθ (x, v, t)a(x, v, t|z)ρt (x, v|z)q(z)dxdvdz

= Eρt (x,v|z)q(z) (aTθ (x, v, t)a(x, v, t|z))

A.5

(40)

Proof of Theorem 4.6

Proof. Consider a given time ti . Based on the expressions for mx , mv , at this moment we exactly have: ρti (x, v|z) = N (x|xi , σx2 I) · N (v|vi , σv2 I) (41) For simplicity, let dsj = dxj dvj Therefore, the Marginal Probability Path is: ρti (x, v) =

Z

ρti (x, v|z)q(z)dz

=

Z

ρti (x, v|(xi , vi ))µ0 (x0 , v0 ) ·

K−1 Y

∗ Kj→j+1 [(xj , vj ), (xj+1 , vj+1 )]ds0 · · · dsK

(42)

j=0

Using the definition of the propagator, we know that: Z

∗ Kj→j+1 [(xj , vj ), (xj+1 , vj+1 )]dxj+1 dvj+1 = 1

(43)

Thus, we can integrate out the future terms (j > i) first , and then iteratively collapse the past terms (j < i):

20

ρti (x, v) =

Z

=

Z

=

Z

=

Z

ρti (x, v|(xi , vi ))µ0 (x0 , v0 ) ·

K−1 Y

∗ Kj→j+1 [(xj , vj ), (xj+1 , vj+1 )]ds0 · · · dsK

j=0

ρti (x, v|(xi , vi ))µ0 (x0 , v0 ) ·

i−1 Y

∗ Kj→j+1 [(xj , vj ), (xj+1 , vj+1 )]ds0 · · · dsi

j=0 ∗ ρti (x, v|(xi , vi ))π0→1 [(x0 , v0 ), (x1 , v1 )]

i−1 Y

∗ Kj→j+1 [(xj , vj ), (xj+1 , vj+1 )]ds0 · · · dsi

j=1

ρti (x, v|(xi , vi ))µ1 (x1 , v1 )

i−1 Y

∗ Kj→j+1 [(xj , vj ), (xj+1 , vj+1 )]ds1 · · · dsi

j=1

.. . =

Z

ρti (x, v|(xi , vi ))µi (xi , vi )dsi

=

Z

(N (x|xi , σx2 I) · N (v|vi , σv2 I))µi (xi , vi )dsi

= (µi ∗ g)(x, v)

(44)

where g(x, v) = (N (x|0, σx2 I) · N (v|0, σv2 I)). When (σx , σv ) → 0, we have ρti (x, v) → µi (x, v) in distribution. If µi (x, v) is absolutely continuous with respect to the Lebesgue measure and admits a density in L1 (R2d ), we can guarantee strong convergence in the L1 norm: limσ→0 ∥ρti (x, v) − µi (x, v)∥L1 = 0 for all i = 0, 1, ..., K.

A.6

Proof of Proposition 4.7

First, we prove that the Marginal Probability Path constructed by TracingFlow converges to the optimal solution of the DOAT problem when σx → 0, σv → 0. We investigate directly in the augmented space. Let s = [x, v]T ∈ X × V, and consider the cost of the DOAT problem. The subsequent position of a particle starting from the initial point s0 can be represented by the flow map ϕt (s0 ). Therefore, the cost can be written as: 1 ∥a(s, t)∥2 ρt (s)dxdvdt 2 t0 X ×V     Z tk Z 1 ∂ϕ −1 ∂ϕ ∥a(ϕt (s0 ), t)∥2 det ρ0 (s0 ) det ds0 dt = ∂s0 ∂s0 t0 X ×V 2 Z Z tk 1 = ρ0 (s0 )ds0 ∥a(ϕt (s0 ), t)∥2 dt 2 X ×V t0  Z Z tk 1 ′′ 2 = ∥γ (t)∥ dt dP[γ] Ω t0 2

LDOAT (a, ρ) =

Z tk Z

(45)

where P is a probability measure in the path space, satisfying the constraint (eti )# P = µi (x, v). Here, et is the evaluation map et [γ] = γ(t), and (eti )# P denotes the push-forward of the probability measure by the evaluation map. According to Proposition 4.1, the minimum cost path connecting si and si+1 is a cubic spline, and the minimum cost is Ci→i+1 [si , si+1 ]. Therefore, for any path γ: Z ti+1 1 ′′ ∥γ (t)∥2 dt ≥ Ci→i+1 [si , si+1 ] (46) 2 ti 21

Thus, the DOAT problem cost has a lower bound: LDOAT (ρ, a) = ≥

Z K−1 X Z ti+1 1 Z K−1 X Ω i=0

≥

ti

Ω i=0

Z K−1 X

2

′′



∥γ (t)∥ dt dP[γ] 2

Ci→i+1 [si , si+1 ]dP[γ]

Ci→i+1 [si , si+1 ]dπ ∗ (si , si+1 )

(47)

i=0

where π ∗ (si , si+1 ) is the Optimal Transport Plan for the SOAT problem between µ(si ) and µ(si+1 ). Consider the path constructed by TracingFlow. In the case where σx , σv → 0, we have ρt (x, v|z) = ρt (s|z) → δ(s − γz∗ (t)), where γz∗ is the optimal single-particle trajectory connecting z = [s0 , s1 , · · · , sK ]. We calculate the transport cost of the constructed path: LTF =

Z

=

Z

q(z)dz q(z)dz

K−1 X Z ti+1 i=0 K−1 X

ti

∗

∥γ ′′ z (t)∥2 dt

Ci→i+1 [si , si+1 ]

i=0

=

Z

µ0 (s0 )

K−1 Y

 ∗ Ki→i+1 [si , si+1 ] 

i=0

K−1 X



Cj→j+1 [sj , sj+1 ] dz

(48)

j=0

Here, using the definition of the propagator: Z

µ0 (s0 ) =

Therefore: LTF =

K−1 Y

∗ Ki→i+1 [si , si+1 ] (Cj→j+1 [sj , sj+1 ]) dz

Z i=0

(49)

∗ πj→j+1 [sj , sj+1 ]Cj→j+1 [sj , sj+1 ]dsj dsj+1

K−1 XZ

∗ πi→i+1 [si , si+1 ]Ci→i+1 [si , si+1 ]dsi dsi+1

(50)

i=0

This coincides exactly with the lower bound of the DOAT problem cost derived above. Thus, the path constructed by TracingFlow is exactly the optimal solution to the DOAT problem in the case where σx , σv → 0. Next, we prove that under the construction of TracingFlow, the acceleration field a(s, t) is single-valued with respect to s and t. By Remark 4.3, under the Optimal Transport Plan, the trajectories of flow maps starting from different points must not intersect. This implies that for a given time t, any point s in the augmented space is traversed by at most one flow map trajectory. Therefore, the acceleration field a(s, t) possesses single-valuedness. In other words, the solution constructed by TracingFlow can indeed be expressed by a single optimal acceleration field a(s, t).

A.7

Proof of Proposition 4.8

Proof. Let γ : [0, T ] → Rd be a path that is C 1 on [0, T ] and piecewise C 2 on each interval [tj−1 , tj ]. The optimal solution is given by a piecewise cubic polynomial: 22

γ(tj−1 + τ ) =

(vj−1 + vj )∆tj − 2(xj − xj−1 ) 3 τ (∆tj )3 3(xj − xj−1 ) − (2vj−1 + vj )∆tj 2 τ + vj−1 τ + xj−1 , + (∆tj )2

(51)

for τ ∈ [0, ∆tj ], where ∆tj = tj − tj−1 , According to Corollary 4.2, it suffices to minimize K X

1 J (v0 , v1 , . . . , vK ) =2 (∥vj−1 ∥2 + ⟨vj−1 , vj ⟩ + ∥vj ∥2 )(∆tj )2 3 (∆t ) j j=1 

+ 3⟨vj−1 + vj , xj−1 − xj ⟩∆tj + 3∥xj−1 − xj ∥

2



(52)

J ≥ 0 always holds, so the the formula above is a semi-positive quadratic form. Take differentiation, we have vj−1 + 2vj 2vj + vj+1 ∇vj J = 2 + ∆tj ∆tj+1

!

xj−1 − xj xj − xj+1 +6 + 2 (∆tj ) (∆tj+1 )2

!

,

for j = 1, . . . , K − 1, (53)

∇v0 J =

2(vK−1 + 2vK ) 6(xK−1 − xK ) 2(2v0 + v1 ) 6(x0 − x1 ) , ∇vK J = + . + 2 ∆t1 (∆t1 ) ∆tK (∆tK )2

(54)

Set the gradients to zero and rearranging the terms to separate the velocities v and positions x, we obtain: For j = 0: 2 1 3(x1 − x0 ) v0 + v1 = . (55) ∆t1 ∆t1 (∆t1 )2 For j = 1, . . . , K − 1: 1 1 1 vj−1 + 2 + ∆tj ∆tj ∆tj+1 For j = K:

!

1 xj − xj−1 xj+1 − xj vj + vj+1 = 3 + ∆tj+1 (∆tj )2 (∆tj+1 )2

1 2 3(xK − xK−1 ) vK−1 + vK = . ∆tK ∆tK (∆tK )2

!

.

(56)

(57)

Denote V = [v0 , v1 , . . . , vK ] ∈ Rd×(K+1) and writing this system in matrix form V A = B, where 2 1  ∆t  1  ∆t1



   0 A=  .  ..    0 

0

2



1 ∆t1 1 1 ∆t1 + ∆t2 1 ∆t2

.. .

0 0



2



0

···

0

0

1 ∆t2

··· .. . .. .

0 .. .

0   ..   . 

1 1 ∆t2 + ∆t3

..

.

··· ···



1 ∆tK−1

0

23



1

2



∆tK−1 1

1 ∆tK−1 + ∆tK 1 ∆tK





 ∈ R(K+1)×(K+1)  0    1  ∆tK  2 ∆tK

(58)

3(x1 − x0 ) x1 − x0 x2 − x1 B= ,3 + ,..., 2 (∆t1 ) (∆t1 )2 (∆t2 )2 "





xK−1 − xK−2 xK − xK−1 3(xK − xK−1 ) 3 + , ∈ Rd×(K+1) 2 2 (∆tK−1 ) (∆tK ) (∆tK )2 

#



(59)

The solution satisfies γ ′ (tj ) = vj for all j = 0, 1, ...K. Since A is strictly diagonally dominant, the solution V is unique.

A.8

Proof of Theorem A.1

Theorem A.1. Let aθ (x, v, t) be Lθ -Lipschitz continuous of (x, v). The Wasserstein-2 distance R (pos) (pos) between the generated distribution ρ̂ti (x) = ρ̂ti (x, v)dv and the true distribution µi (x) in the position space is bounded by (pos)

W22 (ρ̂ti

(pos)

, µi



) ≤ 2eκ(ti −t0 ) Lv0 + (ti − t0 )LAFM |[t0 ,ti ]



(60)

where the cumulative AFM loss and the initial velocity loss are defined as: LAFM |[t0 ,ti ] = Et,(x,v)∼ρt ∥aθ (x, v, t) − a(x, v, t)∥2 , Lv0 = E(x0 ,v0 )∼ρt0 ∥v0,ξ (x0 ) − v0 ∥2 . and κ is a constant only depending on Lθ . Proof. We address the problem directly within the augmented space X × V. Let the system state be denoted by s = [x, v]T , and let µi (s) represent the data distribution in this space. We define a flow map ϕt : [t0 , tK ] × X × V → X × V, which is governed by the exact marginal acceleration field a(s, t) defined in Equation (13). Using (ϕt )# to denote the push-forward operator, the time-dependent density evolves as ρt = (ϕt )# ρt0 , where ρt0 = µ0 . Similarly, let ϕθt be the flow map induced by the learned acceleration field aθ (s, t). The generated distribution ρ̂t then evolves according to ρ̂t = (ϕθt )# ρ̂t0 . It is important to note that ρ̂t0 ̸= ρt0 , as the conditional distribution of the initial velocity is approximated via a neural network. We introduce a bridging measure νt = (ϕθt )# ρt0 , then the Wasserstein-2 distance between joint distributions can be decomposed as 



W22 (ρt , ρ̂t ) ≤ (W2 (ρt , νt ) + W2 (νt , ρ̂t ))2 ≤ 2 W22 (ρt , νt ) + W22 (νt , ρ̂t ) Estimation of Term 1: Define Qt =

(61)

µ0 (s0 )∥ϕθt (s0 )−ϕt (s0 )∥2 ds0 ,here µ0 = ρt0 according " # 0 I to definition. Noting that Qt0 = 0 since the starting points are identical. Define A = . 0 0 Taking the time derivative, we have: R

dQt =2 dt

Z

µ0 (s0 )(ϕθt (s0 ) − ϕt (s0 ))T

dϕθt (s0 ) dϕt (s0 ) − ds0 dt dt

=2

Z

µ0 (s0 )(ϕθt (s0 ) − ϕt (s0 ))T

Aϕθt (s0 ) − Aϕt (s0 ) +

=2

Z

!

"

#

"

#!

0 0 − θ a(ϕt (s0 ), t) aθ (ϕt (s0 ), t)

ds0

"

µ0 (s0 ) (ϕθt (s0 ) − ϕt (s0 ))T A(ϕθt (s0 ) − ϕt (s0 )) "

0 aθ (ϕθt (s0 ), t) − aθ (ϕt (s0 ), t)

"

0 aθ (ϕt (s0 ), t) − a(ϕt (s0 ), t)

+ (ϕθt (s0 ) − ϕt (s0 ))T + (ϕθt (s0 ) − ϕt (s0 ))T

#

##

24

ds0

(62)

We bound the three terms inside the integral separately. Let ∆ϕt (s0 ) = ϕθt (s0 ) − ϕt (s0 ). 1. Since ∥A∥ ≤ 1 , we have ∆ϕt (s0 )T A∆ϕt (s0 ) ≤ ∥∆ϕt (s0 )∥2 . 2. Using the Lipschitz continuity of the network aθ with constant Lθ : "

∆ϕt (s0 )

T

#

0 ≤ ∥∆ϕt (s0 )∥ · ∥aθ (ϕθt (s0 ), t) − aθ (ϕt (s0 ), t)∥ aθ (ϕθt (s0 ), t) − aθ (ϕt (s0 ), t) ≤ Lθ ∥∆ϕt (s0 )∥2 .

(63)

3. Using Cauchy-Schwarz and the inequality 2xy ≤ x2 + y 2 : "

∆ϕt (s0 )

T

#

0 ≤ ∥∆ϕt (s0 )∥ · ∥aθ (ϕt (s0 ), t) − a(ϕt (s0 ), t)∥ aθ (ϕt (s0 ), t) − a(ϕt (s0 ), t) 1 1 ≤ ∥∆ϕt (s0 )∥2 + ∥aθ (ϕt (s0 ), t) − a(ϕt (s0 ), t)∥2 . 2 2 (64)

Substituting these bounds back into the derivative: dQt ≤2 dt

Z

1 1 µ0 (s0 ) (1 + Lθ + )∥∆ϕt (s0 )∥2 + ∥aθ (ϕt (s0 ), t) − a(ϕt (s0 ), t)∥2 ds0 2 2 



= (2Lθ + 3)Qt +

Z

µ0 (s0 )∥aθ (ϕt (s0 ), t) − a(ϕt (s0 ), t)∥2 ds0

(65)

By Gronwall’s inequality, since Qt0 = 0: Qt ≤ e(2Lθ +3)(t−t0 )

Z t Z t0



µ0 (s0 )∥aθ (ϕτ (s0 ), τ ) − a(ϕτ (s0 ), τ )∥2 ds0 dτ

(66)

By the definition of Wasserstein-2 distance, we have: t0 W22 (ρt , νt ) = min π

≤

Z

Z

∥s1 − s2 ∥2 dπ

s.t.

Z

πds1 = ρt ,

Z

πds2 = νt

µ0 (s0 )∥ϕθt (s0 ) − ϕt (s0 )∥2 ds0

= Qt ≤ e(2Lθ +3)(t−t0 )

Z t t0

Es∼µτ ∥aθ (s, τ ) − a(s, τ )∥2 dτ

(67)

= e(2Lθ +3)(t−t0 ) (t − t0 )LAFM |[t0 ,t] .

(68)

Estimation of Term 2:Take π0 (s0 , sξ0 ) as the optimal coupling that attains W22 (νt0 , ρ̂t0 ) = ∥s0 − sξ0 ∥2 dπ0 (s0 , sξ0 ). Note that νt0 = µ0 (the true data distribution). From a pair of starting points (s0 , sξ0 ) ∼ π0 , define the square distance between trajectories under the same approximate flow as ∆(t) = ∥ϕθt (s0 ) − ϕθt (sξ0 )∥2 . Taking the time derivative, we have: R

d∆(t) = 2(ϕθt (s0 ) − ϕθt (sξ0 ))T dt

dϕθt (s0 ) dϕθt (sξ0 ) − dt dt

!

"

= 2(ϕθt (s0 ) − ϕθt (sξ0 ))T

A(ϕθt (s0 ) − ϕθt (sξ0 )) +



#!

0 θ aθ (ϕt (s0 ), t) − aθ (ϕθt (sξ0 ), t)



≤ 2 ∥A∥ · ∥ϕθt (s0 ) − ϕθt (sξ0 )∥2 + ∥ϕθt (s0 ) − ϕθt (sξ0 )∥ · ∥aθ (ϕθt (s0 ), t) − aθ (ϕθt (sξ0 ), t)∥ ≤ 2(1 + Lθ )∆(t).

(69) 25

Again by Gronwall’s inequality, we have: ∆(t) ≤ e2(1+Lθ )(t−t0 ) ∆(t0 ) = e2(1+Lθ )(t−t0 ) ∥s0 − sξ0 ∥2 .

(70)

Since νt = (ϕθt )# νt0 and ρ̂t0 = (ϕθt )# ρ̂t0 , the push-forward of the optimal coupling πt = (ϕθt , ϕθt )# π0 is a valid coupling for (νt , ρ̂t ). Integrate the squared distance over all starting pairs with respect to π0 , we obtain: W22 (νt , ρ̂t ) ≤ ≤

Z

∥ϕθt (s0 ) − ϕθt (sξ0 )∥2 dπ0 (s0 , sξ0 )

Z

e2(1+Lθ )(t−t0 ) ∥s0 − sξ0 ∥2 dπ0 (s0 , sξ0 )

= e2(1+Lθ )(t−t0 ) W22 (νt0 , ρ̂t0 )

(71)

Let π̃(s0 , sξ0 ) be the coupling between ν0 and ρ̂0 induced by the identity mapping on the position space, i.e., pairing s0 = (x0 , v0 ) with sξ0 = (x0 , v0ξ ). We can bound Wasserstein distance by L2 loss: W22 (νt0 , ρ̂t0 ) ≤

Z

∥s0 − sξ0 ∥2 dπ̃(s0 , sξ0 )

=

Z 

=

Z



∥x0 − xξ0 ∥2 + ∥v0 − v0ξ ∥2 dπ̃(x0 , v0 , xξ0 , v0ξ )

∥v0 − v0ξ ∥2 dπ̃(x0 , v0 , xξ0 , v0ξ ) = Lv0

(72)

Thus W22 (νt , ρ̂t ) ≤ e2(1+Lθ )(t−t0 ) Lv0 . Combining the bounds above, we have h

W22 (ρ̂ti , ρti ) ≤ 2 W22 (ρ̂ti , νti ) + W22 (νti , ρti )

i

h

i

≤ 2 e(2Lθ +3)(ti −t0 ) · (ti − t0 )LAFM |[t0 ,ti ] + e(2Lθ +2)(ti −t0 ) Lv0 .

(73)

Define κ = (2Lθ + 3),then the final bound is obtained by: (pos) (pos) (pos) W22 (ρ̂ti , ρti = µi ) = min π

Z

∥x1 − x2 ∥ dπx (x1 , x2 )

≤ min

Z

∥x1 − x2 ∥2 dπ(s1 , s2 )

≤ min

Z

∥s1 − s2 ∥2 dπ(s1 , s2 )

x

π π

2



s.t. s.t.

Z

Z

(pos) πx dx2 = ρti ,

πds2 = ρti ,

Z

Z

(pos)

πx dx1 = ρ̂ti

πds1 = ρ̂ti (74)



= W22 (ρ̂ti , ρti ) ≤ 2eκ(ti −t0 ) (ti − t0 )LAFM |[t0 ,ti ] + Lv0 .

(75)

This completes the proof. Note that our derivation holds even for multi-valued or stochastic cases, since the upper bound on the Wasserstein distance is established via a valid coupling, which remains well-defined regardless of whether the mapping is deterministic.

B

Datasets and Evaluation Metric

B.1

Experiment Setup

The experiments were conducted on a shared high-performance computing (HPC) cluster. All experimental runs utilized GPUs equipped with 80GB of memory. Both the network aθ (x, v, t), employed to approximate acceleration, and the network vϕ (x, v, t), employed to 26

approximate velocity, utilize residual architectures. Specifically, aθ (x, v, t) consists of 4 hidden layers, while vϕ (x, v, t) comprises 2 hidden layers. The Optimal Transport (OT) calculations were implemented using the Python Optimal Transport (POT) library [16], and the training process of neural networks were implemented using Pytorch [43]. In all experiments, TracingFlow employs a completely consistent set of hyperparameters. When solving for OT using the Sinkhorn algorithm, the regularization coefficient is set to ϵ = 1×10−3 . When incorporating biological priors, the penalty coefficient for barcode mismatches is set to p0 = 25. Since the number of cells at each time point in the datasets processed in this paper is on the order of 103 , we do not utilize mini-batch OT; instead, we compute the Transport Plan on the full dataset.

B.2

Evaluation Metrics

We employ the 1-Wasserstein Distance (W1 ) and the 2-Wasserstein Distance (W2 ) to quantify the discrepancy between the predicted distribution and the ground truth distribution. These metrics are defined as follows: W1 (µ, ν) =

W2 (µ, ν) =

Z

min

π∈Π(µ,ν)

min

Z

π∈Π(µ,ν)

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

(76)

!1/2

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

(77)

For the evaluation of datasets containing lineage information, we utilize the lineage-weighted W1 and W2 distances. These are defined based on the decomposition of the distributions. Consider the distribution generated by the model, denoted as µ(x), and the reference distribution provided by the dataset, denoted as ν(x). Both consist of L lineages and can be decomposed as: µ(x) =

L X

wi µ (x), (i)

ν(x) =

i=1

L X

vi ν (i) (x)

(78)

i=1

where µ(i) (x) and ν (i) (x) represent the distributions of the i-th lineage within µ(x) and ν(x), R (i) R (i) respectively, satisfying µ (x)dx = 1 and ν (x)dx = 1. The terms wi and vi denote the P weights of the i-th lineage in µ(x) and ν(x), respectively, subject to the constraint L i=1 wi = PL i=1 vi = 1. The lineage-weighted W1 and W2 distances are defined as: (L)

W1 (µ∥ν) =

L X

vi W1 (µ(i) , ν (i) ),

i=1

(L)

W2 (µ∥ν) =

L X

vi W2 (µ(i) , ν (i) )

(79)

i=1

In essence, we first calculate the W1 and W2 distances between the distributions of corresponding lineages in the two datasets. Subsequently, we compute a weighted average of these distances using the proportions of each lineage given by the reference distribution ν(x). It is important to note that the lineage-weighted W1 and W2 metrics serve solely as indicators of the model’s ability to preserve lineage prior information while learning the dynamics. They are not true distance metrics in the mathematical sense, as they are evidently asymmetric with respect to µ and ν.

B.3

Datasets

3D Simulation Lineage Data. Consider the following dynamics with three lineage barcodes and a twist structure:

27

 (0)  dX = v dt + σx dBtX    t

 (1)  dX = v dt + σx dBtX    t

 (2)  dX = v dt + σx dBtX    t

   

   

   

(0)

dYt

(0)

= σy dBtY

= σz dBtZ

(1)

dYt

(1)

(1)

= −µZt dt + σy dBtY (1)

= µYt

dt + σz dBtZ

(2)

= µZt dt + σy dBtY

(2)

= −µYt

dYt

(2)

(2)

dt + σz dBtZ (80) We take v = 3.0, µ = 3.14, σx = σy = σz = 0.1 and time points t = 0, 1, 2. Initialize (0) (0) (0) (1) (1) (1) (2) (2) (2) (X0 , Y0 , Z0 ) ∼ N ([0, 0, 0], 0.5I), (X0 , Y0 , Z0 ) ∼ N ([0, −2, 0], 0.5I), (X0 , Y0 , Z0 )∼N ([0, 2, 0], 0.5I) and sample 500 particles for each. dZt

dZt

dZt

Figure 5: Dynamics of simulation datasets. (a)2D simulation data, (b)3D simulation data and (c)3D simulation data projected to XY plane. Cite Data. We utilized the Cite-seq dataset introduced in [31], which comprises 31,240 cells collected across 4 distinct time points. To evaluate the performance of TracingFlow on real-world biological data, we applied Principal Component Analysis (PCA) to reduce the dimensionality of the gene expression profiles to 100 and 5 dimensions, respectively. Embryoid Bodies Data. We employed the Embryoid Bodies (EB) dataset from [39], consisting of 16,819 cells sampled at 5 time points. We utilized PCA to reduce the dimensionality of the gene expression data to 5 dimensions. Hematopoiesis Data. We employed the Hematopoiesis dataset from [64], consisting of 130,887 cells sampled at 3 time points. Part of the cells are recorded with barcodes. We selected 22,329 cells with barcodes from the dataset and utilized PCA to reduce the dimensionality of gene expression data to 50 dimensions. When processing this dataset, we filtered for data where barcodes were present at all three time points. The barcodes for the remaining data were set to undefined, incurring no additional transport cost during the incorporation of biological priors (Section 5.3). Mexico Gulf Data To assess the capability and robustness of Tracing Flow in capturing longterm system dynamics, we utilize the Mexico Gulf dataset, following [56]. This two-dimensional dataset comprises 9 time points.

C

Experiment Details

C.1

Experiment Details across various datasets

Simulation 2D Data Table 4 presents the performance of TracingFlow and other algorithms on the 2D Simulation dataset. On average, TracingFlow reconstructs the data distribution at each time point with higher accuracy. Cite 5D & Cite 100D Data We present the distribution reconstruction performance of each algorithm on the Cite 5D and Cite 100D datasets in Table 5 and Table 6, respectively. On these 28

Table 4: W1 and W2 distances between the generated and ground-truth distributions at each time point on the 2D Simulation dataset. Method

t=1

Dynamics

t=2

t=3

t=4

W1

W2

W1

W2

W1

W2

W1

W2

0.1441 ± 0.0164 0.2668 ± 0.1305

0.3644 ± 0.1177 0.5700 ± 0.2574

0.4143 ± 0.1168 0.6108 ± 0.2654

0.5863 ± 0.2615 1.0630 ± 0.2769

0.6238 ± 0.2657 1.3557 ± 0.2770

1.0356 ± 0.2325 2.0986 ± 1.3198

1.1867 ± 0.2890 2.2938 ± 1.3193

1.4459 ± 0.6596 1.9287 ± 0.3002 1.2096 ± 0.0244 0.5117 ± 0.1738 0.1787 ± 0.0258

0.9318 ± 0.2553 1.8387 ± 0.5624 0.9955 ± 0.0921 0.8857 ± 0.3630 0.1753 ± 0.0299

1.0181 ± 0.2050 1.1779 ± 0.5114 1.0009 ± 0.0873 0.8996 ± 0.3689 0.2045 ± 0.0277

1.8079 ± 0.7761 1.8158 ± 0.1794 1.9829 ± 0.1311 1.3274 ± 0.1194 0.4091 ± 0.0523

1.9162 ± 0.7529 1.8740 ± 0.1701 2.0040 ± 0.1582 1.3731 ± 0.1107 0.4821 ± 0.0698

0.5993 ± 0.6176 4.4984 ± 0.2937 1.9740 ± 0.1873 3.9358 ± 0.4572 0.6404 ± 0.2661

0.7265 ± 0.6935 4.6059 ± 0.2977 2.4379 ± 0.2157 3.7880 ± 0.4530 1.1436 ± 0.1577

OT-CFM SF2 M

1st Order

0.1215 ± 0.0164 0.2389 ± 0.1217

3MSBM MMFM HRF CAF TF(Ours)

2nd Order

1.4029 ± 0.6892 1.3226 ± 0.1774 1.2036 ± 0.0425 0.4051 ± 0.1712 0.1609 ± 0.0232

datasets TracingFlow achieves superior reconstruction accuracy at all time steps except for t = 3. The trajectories learned by TracingFlow on these two datasets are illustrated in Figure 6. Table 5: W1 and W2 distances between the generated and ground-truth distributions at each time point on the Cite 5D dataset Method

t=1

Dynamics

OT-CFM SF2 M 3MSBM MMFM HRF CAF TF (Ours)

t=2

t=3

W1

W2

W1

W2

W1

W2

1st Order

0.5615 ± 0.0156 0.5932 ± 0.0249

0.6484 ± 0.0268 0.6797 ± 0.0424

0.8472 ± 0.0463 0.7768 ± 0.0451

0.9416 ± 0.0443 0.8670 ± 0.0480

1.0536 ± 0.0534 0.9251 ± 0.0279

1.1566 ± 0.0573 1.0221 ± 0.0352

2nd Order

5.3735 ± 0.3961 0.5335 ± 0.0134 0.4225 ± 0.0152 5.1724± 0.0452 0.2682 ± 0.0241

5.4244 ± 0.3582 0.5911 ± 0.0057 0.4622 ± 0.0185 5.6242± 0.0763 0.2916± 0.0312

4.3118 ± 0.4589 2.2290 ± 0.2151 0.6227 ± 0.0421 5.4333 ± 0.0914 0.6094 ± 0.0615

4.5458 ± 0.4141 2.4557 ± 0.1613 0.6814 ± 0.0392 5.4637 ± 0.1425 0.6516 ± 0.0684

1.2254 ± 0.2998 4.8413 ± 0.2538 0.7638 ± 0.0738 3.0710 ± 0.1853 0.8808 ± 0.1208

1.3427 ± 0.3307 5.2261 ± 0.4061 0.8394 ± 0.0713 3.1354 ± 0.2672 0.9725 ± 0.1329

Table 6: W1 and W2 distances between the generated and ground-truth distributions at each time point on the Cite 100D dataset Method

t=1

Dynamics

OT-CFM SF2 M 3MSBM MMFM HRF CAF TF (Ours)

t=2

t=3

W1

W2

W1

W2

W1

W2

1st Order

10.1136 ± 0.0104 10.3598 ± 0.0051

10.1873 ± 0.0135 10.4448 ± 0.0053

10.2504 ± 0.0209 11.4647 ± 0.0384

10.3145 ± 0.0201 11.5456 ± 0.0402

11.2682 ± 0.0938 15.2032 ± 0.1813

11.3642 ± 0.0997 15.4367 ± 0.1949

2nd Order

19.9677± 0.1373 9.8726 ± 0.0383 9.9660 ± 0.0215 18.1182 ± 0.3214 2.8955 ± 0.1429

20.1015 ± 0.1105 9.9305 ± 0.0361 10.0580 ± 0.0341 18.1971 ± 0.2845 3.2763 ± 0.1776

15.3218 ± 0.1232 12.1374 ± 0.2422 9.9451 ± 0.0632 17.4947 ± 0.4267 9.8235 ± 0.1489

15.4089 ± 0.1343 12.3076 ± 0.2495 9.9986 ± 0.1526 17.5767 ± 0.4912 9.8773± 0.1827

12.0229 ± 0.2895 15.5960 ± 1.4972 10.5485 ± 0.1082 14.6158 ± 0.5632 11.4910 ± 0.1693

12.1395 ± 0.3621 16.1222 ± 2.0195 10.6442 ± 0.1254 15.1773 ± 0.5215 11.6054± 0.2064

EB 5D Data Table 7 displays the interpolation accuracy on the EB 5D dataset, evaluated by holding out one time point during training. TracingFlow outperformed other methods in reconstruction accuracy for all time points except t = 1. Figure 7 illustrates the interpolated distributions generated by TracingFlow across the four time points. Table 7: W1 and W2 distances between predicted and ground-truth distributions at held-out time points on the EB 5D dataset. Method

t=1

Dynamics

t=2

t=3

t=4

W1

W2

W1

W2

W1

W2

W1

W2

3.8462 ± 0.1449 3.6713 ± 0.0589

4.8404 ± 0.1496 3.5511 ± 0.0394

5.5792 ± 0.1358 3.9441 ± 0.0418

5.0713 ± 0.2839 4.3175 ± 0.2517

5.6448 ± 0.2861 4.8738 ± 0.3247

7.1368 ± 0.3799 6.9516 ± 0.0616

7.6938 ± 0.4574 7.4727 ± 0.0795

6.5434 ± 1.2227 11.8376 ± 0.7823 3.8930 ± 0.0243

7.3626 ± 0.8043 5.7932 ± 0.2162 3.5002 ± 0.0983

8.0596 ± 1.0252 6.2242 ± 0.1828 3.8035 ± 0.1086

7.5862 ± 1.1170 9.3326 ± 0.5617 3.5643 ± 0.1140

8.1734 ± 0.7900 9.9837 ± 0.4994 3.7947 ± 0.4123

7.7029 ± 0.4241 13.2046 ± 2.1528 7.3639 ± 0.1316

7.1022 ± 0.3971 14.3464 ± 2.2982 7.6523 ± 0.1976

OT-CFM SF2 M

1st Order

3.4335 ± 0.1465 3.4204 ± 0.0637

3MSBM MMFM TF (Ours)

2nd Order

6.1086 ± 1.1215 10.7803 ± 0.6718 3.5611 ± 0.0362

3D Simulation Lineage Dataset Table 8 demonstrates the ability of various algorithms to preserve biological priors on the 3D Simulation Lineage dataset. Since the trajectories learned 29

a) CITE 5D

b) CITE 100D

Figure 6: Trajectories learned by TracingFlow on the (a) Cite 5D and (b) Cite 100D datasets, projected to 2D using PCA. a) EB5D w/o t=1

b) EB5D w/o t=2

c) EB5D w/o t=3

d) EB5D w/o t=4

Figure 7: Predicted distributions by TracingFlow at four held-out time points, projected to 2D using PCA. by TracingFlow align with biological priors (see Figure 3), it achieves lower lineage-weighted W1 and W2 scores on average. Table 8: lineage-weighted W1 and W2 distances between the generated and ground-truth distributions at each time point on the 3D Simulation Lineage dataset Method OT-CFM SF2 M 3MSBM MMFM HRF CAF TF w/o Bio. Prior TF(Ours)

t=1

Dynamics

t=2

W1

W2

W1

W2

1st Order

2.6080 ± 0.0029 2.8262 ± 0.0409

2.6869 ± 0.0030 2.8373 ± 0.0462

1.8686 ± 0.0382 0.2247 ± 0.0617

1.9456 ± 0.0403 0.2533 ± 0.0568

2nd Order

4.0141 ± 0.8788 1.8585 ± 0.0350 2.9150 ± 0.2595 1.6284 ± 0.2249 2.8267 ± 0.2896 0.4282 ± 0.1606

4.2262 ± 0.8596 2.1618 ± 0.0496 2.9265 ± 0.3901 1.6623 ± 0.3726 2.8376 ± 0.3714 0.4989 ± 0.1585

0.7548 ± 0.0726 3.5695 ± 0.5692 0.6717 ± 0.5491 3.7328 ± 0.3305 0.2317 ± 0.0512 0.4795 ± 0.2680

0.6751 ± 0.0573 4.1028 ± 0.6037 0.6905 ± 0.4143 3.8471 ± 0.3316 0.2687 ± 0.0892 0.5667 ± 0.2654

Hematopoiesis Dataset Table 9 presents the performance of different algorithms in retaining biological priors on real-world datasets, measured by the lineage-weighted W1, W2 distance. TracingFlow surpassed other algorithms at all time points. We also evaluated the algorithms 30

without introducing biological priors by Vanilla W1 , W2 distance, and as shown in Table 10, TracingFlow maintained superior performance. Figure 8 visualizes the data points generated by TracingFlow on the real-world dataset under both conditions (with and without biological priors). Table 9: lineage-weighted W1 and W2 distances between the generated and ground-truth distributions at each time point on the Hematopoiesis dataset Method

t=1

Dynamics

OT-CFM SF2 M 3MSBM MMFM HRF CAF TF w/o Bio. Prior TF

t=2

W1

W2

W1

W2

1st Order

13.2972 ± 0.0804 14.1427 ± 0.0341

13.7134 ± 0.0831 14.4199 ± 0.0303

16.6918 ± 0.1973 23.2913 ± 0.2590

17.0807 ± 0.2025 23.5314 ± 0.2640

2nd Order

17.6487 ± 0.5250 18.4040 ± 2.8129 14.1007 ± 0.0804 13.1098 ± 0.0628 14.1260 ± 0.0526 11.5203 ± 0.0695

17.8710 ± 0.5204 19.8538 ± 2.8906 14.3122 ± 0.0831 13.3263 ± 0.0742 14.5958 ± 0.0537 12.0151 ± 0.0809

19.0139 ± 1.4418 36.0412 ± 3.0429 16.8135 ± 0.1973 17.3232 ± 0.0580 19.1073 ± 0.1932 16.5736 ± 0.2603

19.4144 ± 1.3622 40.2282 ± 3.3934 17.1154 ± 0.2025 17.6685 ± 0.0572 19.4262 ± 0.2375 16.9027 ± 0.2999

Table 10: W1 and W2 distances between the generated and ground-truth distributions at each time point on the Hematopoiesis dataset Method OT-CFM SF2 M 3MSBM MMFM HRF CAF TF w/o Bio. Prior

a)

t=1

Dynamics

t=2

Average

W1

W2

W1

W2

W1

W2

First-order

9.8780 10.2352

10.6132 10.8017

10.7621 13.2507

11.4466 13.6879

10.3201 11.7430

11.0299 12.2428

Second-order

12.8214 10.4305 10.6233 21.8716 7.0139

13.2501 10.8695 11.4510 22.3309 7.4140

20.0262 22.3269 10.9570 17.1251 10.1642

20.6374 23.5322 11.6201 18.3914 10.5925

16.4238 16.3787 10.7902 19.4984 8.5891

16.9438 17.2008 11.5356 20.3612 9.0032

c)

b)

Figure 8: Original data (a), and data points generated by TracingFlow on the Hematopoiesis dataset without biological priors (b) and with biological priors (c), projected to 2D using UMAP. Mexico Gulf Dataset Table 11 presents the performance of various algorithms on the Mexico Gulf dataset. Tracing Flow achieves the best distribution reconstruction accuracy at all time points, yielding smooth trajectories.

31

Table 11: W1 and W2 distances between the generated and ground-truth distributions at each time point on the Mexico Gulf dataset Method OT-CFM SF2 M 3MSBM MMFM HRF CAF TF Method OT-CFM SF2 M 3MSBM MMFM HRF CAF TF

C.2

t=1

Dynamics

t=2

t=3

t=4

W1

W2

W1

W2

W1

W2

W1

W2

First-order

0.0336 0.0768

0.0375 0.0917

0.0791 0.0802

0.0887 0.0964

0.1367 0.0819

0.1415 0.0968

0.1409 0.0809

0.1474 0.0987

Second-order

0.4407 0.0346 0.6205 1.7659 0.0128

0.6305 0.0364 0.6217 1.7663 0.0145

0.5565 0.2961 1.2004 2.4713 0.0301

0.7617 0.3008 1.2011 2.4714 0.0339

0.5887 1.0538 1.5334 2.0742 0.0407

0.9113 1.0571 1.5337 2.0755 0.0458

0.5231 1.8755 1.4986 1.3040 0.0348

1.0714 1.8766 1.5013 1.3071 0.0456

t=5

Dynamics

t=6

t=7

t=8

Average

W1

W2

W1

W2

W1

W2

W1

W2

W1

W2

First-order

0.1318 0.1030

0.1370 0.1191

0.1364 0.0851

0.1425 0.0956

0.1610 0.1401

0.1799 0.1494

0.1871 0.1965

0.2236 0.2068

0.1258 0.1056

0.1373 0.1193

Second-order

0.2185 2.3926 1.3022 0.7086 0.0545

1.2783 2.3928 1.3075 0.7124 0.0657

0.5362 2.7214 0.9476 0.4669 0.0778

2.0436 2.7215 0.9518 0.4751 0.0824

0.8474 2.9191 0.6480 0.6672 0.1039

4.5953 2.9199 0.6535 0.6914 0.1133

1.3606 2.9930 0.3296 0.9641 0.1206

13.2434 2.9945 0.3341 0.9914 0.1309

0.6340 1.7857 1.0101 1.3028 0.0594

3.0669 1.7875 1.0131 1.3113 0.0665

Training time and scalability of TracingFlow

To evaluate the scalability of TracingFlow, we report the training time required by each method across various datasets, as presented in Table 12. Experimental results indicate that Flow Matching algorithms based on second-order dynamics (e.g., 3 MSBM and our TracingFlow) generally require more training time compared to first-order dynamics algorithms (e.g., OT-CFM and SF2 M). Notably, the training time of TracingFlow is slightly lower than that of 3 MSBM, demonstrating that our method does not introduce additional computational overhead. Table 12: Training time comparison (in seconds) of different methods across various datasets.

C.3

Method

Simulation

Cite 5D

Cite 100D

SimLineage-3D

RealLineage

OT-CFM SF2 M MMFM 3 MSBM CAF HRF TF

12 s 26 s 42 s 220 s 94 s 247 s 225 s

175 s 79 s 96 s 495 s 274 s 271 s 319 s

218 s 84 s 179 s 910 s 382 s 359 s 806 s

13 s 34 s 68 s 244 s 256 s 289 s 269 s

196 s 61 s 201 s 855 s 264 s 347 s 742 s

Sensitivity Analysis on Minibatch-OT

In all experiments presented in this paper, the full dataset was utilized to compute the OT plan. However, as the scale of the data increases, the computational cost of the OT plan grows rapidly, making the adoption of the minibatch computation method described in Section 5.2 inevitable. We conducted a sensitivity analysis on the Cite-5D dataset to evaluate the impact of batch size in minibatch-OT computation on the accuracy of distribution reconstruction. The results are presented in Table 13. Experimental results demonstrate that TracingFlow consistently and accurately reconstructs the distributions at each time point when the batch size is adjusted within a certain range ( batch size ∼ 103 ).

32

Table 13: Sensitivity analysis of batch size for Minibatch-OT on the Cite-5D dataset. t=1

Batch Size 1000 2000 Full Dataset

t=2

t=3

Average

W1

W2

W1

W2

W1

W2

W1

W2

0.4344 0.4196 0.2868

0.4798 0.4624 0.3185

0.6522 0.5640 0.5487

0.7183 0.6251 0.5997

0.7658 0.7972 0.7689

0.8479 0.8820 0.8499

0.6175 0.5936 0.5348

0.6820 0.6565 0.5894

33

D

Pseudocode for Algorithm

The pseudocode for the training process of Tracing Flow is provided in Algorithm 1. Algorithm 1 Tracing-Flow Matching Training (i)

k Require: Snapshot datasets Dk = {xk }N i=1 at times t0 < t1 < · · · < tK ; Acceleration Network aθ ; Velocity Network vξ ; Hyperparameters: batch size Btraj , learning rates ηa , ηv . (i) Nk ← {0} for all k. 1: Initialize: Estimated velocity sets Vk = {vk }i=1 2: while Velocity Estimates Not Converged do 3: Stage 1: Velocity Assignment (Iterative SOAT) (i) 4: Initialize velocity accumulators V̂k ← ∅ for all k, i. 5: // Step 1.1: Solve Optimal Transport 6: for k = 0 to K − 1 do 7: Compute cost matrix Ci→i+1 with / w.o. biological prior. ∗ 8: Compute optimal coupling πk→k+1 between (xk , vk ) and (xk+1 , vk+1 ) by solving the SOAT Problem. ∗ ∗ 9: Compute propagator (transition) matrix Kk→k+1 by normalizing πk→k+1 :

10:

π∗

j′

11: 12: 13: 14: 15: 16: 17: 18: 19: 20: 21: 22: 23:

(i,j) (i,j ′ ) k→k+1

∗ Kk→k+1 (i, j) = P k→k+1 π∗

end for // Step 1.2: Trajectory Sampling & Velocity Update (s ) (s ) Sample Btraj discrete trajectories Tm = {(xk m,k , vk m,k )}K k=0 for m = 1 . . . Btraj . · Initial indices sm,0 are sampled uniformly. ∗ · Subsequent indices sm,k+1 ∼ Kk→k+1 (sm,k , ·). for m = 1 to Btraj do (s

)

(s

)

Compute optimal path velocities {v̂k m,k }K k=0 for trajectory Tm using Proposition 4.8. (s

)

(s

)

Accumulate: V̂k m,k ← V̂k m,k ∪ {v̂k m,k } for k = 0 . . . K. end for // Step 1.3: Average Accumulated Velocities for k = 0 to K, i = 1 to Nk do (i) if V̂k ̸= ∅ then P (i) vk ← 1(i) (i) v v∈V̂ |V̂k |

k

end if 25: end for 26: end while 27: Stage 2: Final Transport Plan Computation 28: Compute optimal couplings π ∗ and propagators K∗ using the converged velocities Vk . 29: Stage 3: Neural Network Training 30: while Not Converged do 31: Sample random time t ∼ U [t0 , tK ] and coupling z ∼ q(z). 32: Compute Acceleration Matching Loss: h i 33: LCAFM = E(x,v)∼ρt (·|z),z∼q(z),t∼U [t0 ,tK ] ∥aθ (x, v, t) − a(x, v, t|z)∥2 34: Compute Initial Velocity Loss: i h 35: Lv0 = E(x,v)∼ρ0 ∥vξ (x) − v∥2 36: Update parameters (Gradient Descent): 37: θ ← θ − ηa ∇θ LCAFM 38: ξ ← ξ − ηv ∇ξ Lv0 39: end while 24:

34

E

Discussion

E.1

Relation to other Algorithms

Relation to 3 MSBM 3 MSBM considers the following stochastic optimal control problem: min Eρt

Z

[∥at ∥2 ]dt

dmt = Amt dt + ut dt + gdWt , "

#

"

xn ∼ qn =

(81) Z

πn (x, v)dvn

(82)

#

0 0 0 0 where mt = [xt , vt ]T , A = , and g = . This problem can be viewed as a “stochastic” 0 1 0 σ version of the VM-DOAT problem proposed in this paper. 3 MSBM also provides a simulationfree training method to regress the acceleration field. However, 3 MSBM does not explicitly solve this stochastic optimal control problem using the form of optimal coupling plus optimal single-particle trajectories. Its optimal coupling requires iterative solutions similar to Rectified Flow [36] : a process of repeatedly selecting the coupling and training aθ . Furthermore, it relies on a heuristic method to estimate the initial velocity (via normal initialization and iteratively solving forward and backward SDEs), rather than incorporating this estimation as part of the optimal control problem. Experiments show that Tracing Flow achieves better distribution reconstruction accuracy than 3MSBM. Relation to MMFM

In MMFM, the single-particle path is similarly specified by: Z

∥γ ′′ (t)∥2 dt,

xk = γ(tk )

(83)

which corresponds to a natural spline. However, MMFM does not solve an optimal control problem over distributions. Furthermore, the velocity field in MMFM remains single-valued with respect to x, which prevents it from learning trajectories that cross in the position space. Relation to OAT-FM OAT-FM introduces a DOAT problem formulation consistent with this paper; however, the objective of OAT-FM is not to genuinely solve the DOAT problem within the augmented space X × V, but rather to fine-tune a pre-trained FM model. It employs the loss: LCAFM = Eπ∗ [(x0 ,v0 ),(x1 ,v1 )] α∥ 

+ α∥

xt − x0 v0 + vθ 2 − ∥ + (1 − α)∥vθ (xt , t) − v0 ∥2 t 2

 x1 − xt vθ + v1 2 − ∥ + (1 − α)∥v1 − vθ (xt , t)∥2 1−t 2

(84)

to enforce the velocity field along the line segment connecting any two samples to be as parallel to that line as possible. As a fine-tuning approach, it effectively still addresses a distribution transport problem involving only two time points within the position space X , rather than tackling a multi-marginal optimal control problem. Furthermore, OAT-FM does not establish a 1 R1 connection between its loss function and the cost along each trajectory ∥γ̈∥2 dt, nor does it 2 0 explicitly learn the acceleration field a(x, v, t).

E.2

Is S = X × V a Phase Space?

In the Conclusion and Limitation section, we noted that if x is regarded as the generalized R coordinate and v as its time derivative, the term 12 ∥a∥2 dt cannot be interpreted as a physical action, as the action in classical mechanics typically does not involve second-order time derivatives. Here, we offer an alternative physical interpretation of this control objective. 35

Proposition E.1. Consider a d-dimensional second-order dynamical system governed by the equations: ẋ = v

(85)

v̇ = a

(86)

where x, v, a ∈ Rd . The optimal control objective is given by: 1 min 2

Z T 0

∥a∥2 dt

(87)

The evolution of this system under the optimal control law is equivalent to a Hamiltonian system in 4d-dimensional phase space with the following Hamiltonian: 1 H = pT1 q2 + ∥p2 ∥2 2

(88)

where q1 , q2 ∈ Rd are the canonical coordinates, and p1 , p2 ∈ Rd are the canonical momenta (note that q1 does not appear in the Hamiltonian). Here, q1 and q2 correspond to x and v in the optimal control problem, respectively, while p2 corresponds to a. Proof. We derive the equations of motion directly using Hamilton’s canonical equations: q̇1 =

∂H = q2 , ∂p1

q̇2 =

∂H = p2 , ∂p2

ṗ1 = −

∂H = 0, ∂q1

ṗ2 = −

∂H = −p1 ∂q2

(89)

This implies that p1 is constant. Solving these equations sequentially yields: p2 (t) = c′2 − p1 t 1 q2 (t) = c′1 + c′2 t − p1 t2 2 1 1 q1 (t) = c′0 + c′1 + c′2 t2 − p1 t3 2 6

(90) (91) (92)

Comparing this with the results in Equation (9), we observe that q1 (t) and q2 (t) satisfy the same equations as x(t) and v(t), with the undetermined constants determined by the boundary conditions. This completes the proof. Consequently, we observe that although the augmented space S = X × V concatenates position x and velocity (momentum) v, it cannot be interpreted as a phase space. Instead, it should be viewed as the configuration space of the classical mechanical system described by the Hamiltonian above. We can attempt to recover the Lagrangian from this Hamiltonian. The canonical equations established that q̇1 = q2 and q̇2 = p2 . Applying the Legendre transformation: L = pT1 q̇1 + pT2 q̇2 − H 1 = pT1 (q̇1 − q2 ) + ∥q̇2 ∥2 2

(93)

We find that the term pT1 cannot be eliminated; the Lagrangian cannot be expressed solely as a function of the canonical coordinates and their first derivatives. This is a characteristic of constrained systems (in Dirac’s theory of constraints, a constrained system is typically defined as one where the Hessian matrix of the Lagrangian is singular. To see this intuitively, if we regard p1 in the Lagrangian as d new canonical coordinates, their corresponding canonical momenta are zero, implying the motion is confined to a submanifold of the phase space). In Equation (10), we derived the minimum cost for the system to transition from (x1 , v1 ) to (x2 , v2 ). While this cost is a symmetric positive-definite quadratic form with respect to x 36

and v, it does not constitute a metric on the augmented space S = X × V. (If it were a metric, the augmented space would be flat, and the optimal control trajectories—geodesics—would be linear functions of t). Maupertuis’ principle suggests a close relationship between the metric in the configuration space and the system’s Lagrangian [3]: if the Lagrangian takes the form: 1 L(q, q̇) = mij (q)q̇ i q̇ j − V (q) 2

(94)

then the configuration space possesses the metric: gij (q) = 2(E − V (q))mij (q)

(95)

where E is the total energy of the particle. For the optimal control problem discussed herein, we have seen that the Lagrangian cannot be written in such a form; therefore, it is fundamentally impossible to equip the configuration space S = X × V with such a metric.

E.3

Broader Impacts

This paper presents work whose goal is to advance the field of machine learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here. However, as our algorithm is applied to single-cell data analysis, the fidelity of the generated trajectories is heavily dependent on the quality of the input data. Consequently, when utilizing our method for biological or medical research purposes, it is critical to employ high-quality datasets, incorporate established and accurate biological priors, and ensure that the algorithmic outputs are rigorously validated by domain experts.

37

Related documents

• TracingFlow: A Simulation-Free Trajectory Inference Framework Based on Second-Order Dynamics#662245
• TracingFlow: A Simulation-Free Trajectory Inference Framework Based on Second-Order Dynamics#661195
• One Diffusion Model, Two Roles: Guided Trajectory Planning and Safety-Critical Scenario Generation in Closed-Loop Simulation#971578
• EditWM: Event-Decomposed World Modeling with Incremental Correction for End-to-End Autonomous Driving#1037000
• FUSE: FK-Steered Multi-Modal Flow Matching for Efficient Simulation-Based Posterior Estimation#343481