arXiv:2605.18587v1 [q-bio.GN] 18 May 2026
PACE: Geometry-Aware Bridge Transport for Single-Cell Trajectory Inference
Chenglei Yu1,2∗, Chuanrui Wang2∗, Bangyan Liao1,2 & Tailin Wu2† 1 Zhejiang University 2 Department of Artificial Intelligence, School of Engineering, Westlake University {yuchenglei, wangchuanrui, wutailin}@westlake.edu.cn
Abstract Single-cell trajectory inference from destructive time-course snapshots is fundamentally ill-posed: neither cross-time cell correspondences nor the continuous paths between snapshots are observed, so the observed snapshot distributions alone do not uniquely determine the underlying dynamics. Existing optimal transport and flow-based methods typically couple cells by Euclidean proximity at observed clock times, which can misalign trajectories when development is asynchronous and cells sampled at the same experimental time occupy different latent pseudotime stages. We propose PACE, a trajectory inference framework that selects geometry-consistent continuous transport dynamics from destructive time-course snapshots through three coupled components. First, PACE constructs a state- and time-dependent anisotropic Riemannian metric that preserves low cost along locally supported tangent directions while penalizing normal velocity components. Second, it alternates between refining cross-time couplings under the induced path-action cost and fitting endpoint-preserving neural bridges between adjacent snapshots. Third, it distills the learned bridge dynamics into a global continuous-time velocity field over cellular states. Across seven controlled and biological datasets covering nine held-out reconstruction experiments, PACE achieves the strongest overall reconstruction performance, reducing MMD, W1 , and W2 by 23.7% on average relative to the strongest competing baseline. PACE also improves RNA-velocity alignment by 15.4% on an embryoid body differentiation benchmark, without requiring explicit cell pairing, lineage tracing, or RNA velocity supervision during training. Code is available at https://github.com/AI4Science-WestlakeU/PACE.
1
Introduction
Understanding how cells move from one state to another is a central problem in single-cell biology [1, 2, 3, 4, 5]. Processes such as development, differentiation, immune activation, tumor evolution, and cellular reprogramming are not simply collections of discrete cell types, but continuous populationlevel transitions through high-dimensional molecular state space [6, 7, 8, 9]. The key questions are therefore dynamical, including which early states commit to particular fates, which intermediate states are transient but decisive, and which regulatory programs drive these transitions. Trajectory inference [10, 11] aims to answer these questions by reconstructing continuous cell-state evolution from time-course single-cell observations. Transport- and flow-based trajectory methods provide a natural framework for modeling such dynamics [3, 12, 10]. Since time-course single-cell data provide population-level snapshots rather than paired observations of the same cells [3, 13], these methods must infer how marginal distributions are coupled across time. Most existing formulations operate in the Euclidean representation space, ∗ Equal contribution.
Preprint.
†Corresponding author.
where couplings are estimated from Euclidean endpoint costs and continuous paths or velocity fields are regularized by Euclidean kinetic energy, action, or smoothness [14, 15, 16, 17]. This Euclidean geometry is convenient, but it is not necessarily aligned with biological progress. Two cells can be close in expression space while lying at different developmental stages or fate branches, whereas biologically plausible motion may follow curved and locally anisotropic directions along the cell-state manifold [18, 19]. This mismatch makes trajectory inference from destructive snapshots fundamentally ambiguous [13, 20]: the cross-time coupling is not observed and is generally not identifiable from marginal distributions alone. A coupling may satisfy the marginal constraints while connecting cells across incompatible developmental programs, stalled states, or fate branches. For example, in mouse reprogramming [3], differences in growth rates between cell types can make marginal matching misleading, coupling apoptotic stromal cells to rapidly expanding iPSCs. Human iPSC reprogramming [21] shows the same issue at finer granularity, that cells collected on the same day can contain heterogeneous primed-like, naive-like, and trophectoderm-like intermediates, so clock-time adjacency alone does not determine which cells should be coupled. Auxiliary measurements such as RNA velocity, lineage tracing, or metabolic labeling [22, 23, 24] can help, but require additional experiments and are unavailable in many datasets. Existing Euclidean or support-based regularizers can encourage short, smooth, or data-supported paths [25], but they do not directly encode which local directions of motion are developmentally admissible. We propose PACE, a geometry-aware trajectory-inference framework for selecting snapshot-consistent continuous transport dynamics from unpaired time-course snapshots. The intuition is that, although true cell identities and intermediate paths are unobserved, local spatiotemporal neighborhoods still provide a weak but useful prior on admissible directions, since motion along locally supported tangent directions is more plausible than motion orthogonal to the observed manifold structure. PACE therefore replaces Euclidean transport costs with a Riemannian path action induced by a state- and time-dependent anisotropic metric, which assigns lower cost to locally supported tangent motion and higher cost to normal motion. Specifically, PACE estimates local tangent directions from spatiotemporal neighborhoods of observed cells and uses this metric to define an anisotropic bridge transport problem, where the cost of coupling two cells is the minimum Riemannian action of an endpoint-conditioned path. PACE then alternates between refining OT couplings under this path-action cost and fitting endpoint-preserving neural bridges, allowing local geometry to influence both which cells are coupled and how they move between snapshots. Finally, the learned bridge dynamics are distilled into a continuous-time population velocity field over cellular states. Our contributions are threefold: 1. We introduce a time- and state-dependent Riemannian metric based on local spatiotemporal tangent-space projections, providing a geometry-aware path-action cost for selecting couplings between adjacent destructive snapshots beyond Euclidean endpoint proximity. 2. We develop an iterative bridge-transport procedure that alternates between Riemannian coupling refinement and endpoint-preserving neural bridge fitting, allowing cross-time correspondences and interpolant paths to be selected jointly under the same geometry-aware action. 3. We distill endpoint-conditioned bridges into a global continuous-time population velocity field and evaluate PACE on seven datasets covering nine held-out reconstruction experiments. PACE achieves the strongest overall reconstruction performance, reducing MMD, W1 , and W2 by 23.7% on average relative to the strongest competing baseline, improves RNAvelocity alignment on an embryoid body differentiation benchmark, and shows consistent gains across component ablations.
2
Related Work
Optimal transport and flow-based trajectory inference. Optimal transport provides a natural framework for coupling unpaired snapshot distributions in single-cell analysis [3, 10]. Recent neural extensions such as TrajectoryNet [16], Wasserstein Lagrangian Flows [14], and Conditional Flow Matching (CFM) [15, 26] learn continuous dynamics by regressing vector fields along interpolants between coupled endpoints. Schrodinger bridge methods [27, 28, 29] add stochasticity or entropic regularization to the transport problem, while Curly Flow Matching [30] extends the framework 2
| Tangent:
Time:
Normal:
(b) Local Anisotropic Metric
(a) Unpaired Snapshots
(c) Geometry-aware Bridge Transport
Local PCA Induced Normal Projector
Objective
Metric at
Tangent direction (low cost)
Path Cost
Normal direction (high cost)
Stage 1: Bridge Update Poor Initial Paths
Stage 3: Velocity Field Learning
Stage 2: Coupling Refinement
Geometry-aligned Paths
Incorrect Coupling
Refined Coupling
Update rule Update rule
Velocity distillation
Iterative Bridge / Coupling Update
Figure 1: Overview of PACE. PACE uses local PCA to construct an anisotropic metric Gk (x, t) = (k) I + αCN (x, t), trains endpoint-preserving neural bridges under the corresponding path-action cost, iteratively refines cross-time couplings, and distills the learned bridge dynamics into a global velocity field for trajectory inference from unpaired snapshots. to non-gradient dynamics using approximate velocity information. These methods typically infer couplings from Euclidean proximity or OT distances in the observed space, which can misalign trajectories when cells at the same experimental time occupy different pseudotime stages. Geometry-aware generative models. The manifold hypothesis has motivated data-dependent Riemannian metrics in ambient spaces [31, 32], flow matching on known manifolds or general geometries [33], and manifold-aware OT flows for single-cell trajectories [17]. Metric Flow Matching (MFM) [34] is closest to our setting: it uses task-independent support-aware metrics such as LAND [32] and RBF to pull geodesics toward the data support. PACE differs in both the source and use of geometry. In asynchronous reprogramming, density support alone does not determine plausible cross-time transitions, since a high-density intermediate may connect to either progressed or refractory fates and Euclidean proximity can couple incompatible programs. PACE instead builds a time- and state-dependent, direction-aware metric from local spatiotemporal tangent subspaces estimated from destructive snapshots, penalizing motion orthogonal to plausible developmental directions. The resulting metric action is used both to learn interpolant paths and to refine cross-time couplings, whereas MFM typically assumes the endpoint pairing is fixed before interpolant learning.
3
Method
3.1
Problem formulation
We observe single-cell point clouds at anchor times {tk }K k=0 , with t0 = 0, (k)
d k X (k) = {xi }N i=1 ∼ ρ̂k ⊂ R ,
k = 0, 1, . . . , K.
Here d denotes the representation dimension. Our goal is to infer trajectories and a velocity field transporting ρ̂k to ρ̂k+1 between adjacent times. Once learned, the velocity field can be integrated from xt0 ∼ ρ̂0 to generate trajectories over the time course. Proposition 1 (Ill-posedness from snapshots alone). Given only the point clouds {X (k) }K k=0 , the reconstruction of cross-time couplings and intermediate trajectories is non-unique. Indeed, for a single interval [tk , tk+1 ], many couplings πk ∈ Π(ρ̂k , ρ̂k+1 ) share the same endpoint marginals. For any such coupling, each coupled endpoint pair can also be connected by infinitely many smooth paths. Thus, the observed snapshots determine neither a unique correspondence nor a unique intermediate trajectory. This observation motivates PACE to frame trajectory inference as a 3
variational problem over couplings and paths, selected by a geometry-aware path action. The proof of Proposition 1 is given in Appendix D. 3.2
Geometry-aware bridge transport
To address the ill-posedness in Proposition 1, PACE formulates trajectory reconstruction as selecting both a cross-time coupling and a family of continuous paths, chosen by minimizing a geometry-aware action. Given a time- and state-dependent metric tensor Gk (x, t) (constructed in §3.3), the cost of a path γ, or "path action", between endpoints (x, y) is Z 1 AGk [γ] =
γ̇(τ )⊤ Gk γ(τ ), tk,τ γ̇(τ ) dτ,
γ(0) = x, γ(1) = y,
(1)
0
where τ ∈ [0, 1] is local interpolation time and the physical time is tk,τ = tk + τ (tk+1 − tk ). The path action first defines the geometry-aware endpoint cost, and the ideal PACE bridge problem then transports mass using this induced cost: cGk (x, y) = inf AGk [γ], γ(0)=x γ(1)=y
πk⋆ ∈ arg
Z min
πk ∈Π(ρ̂k ,ρ̂k+1 )
cGk (x, y) dπk (x, y).
(2)
Here cGk (x, y) is not a fixed Euclidean endpoint distance; it is the minimum action required to move from x to y under the metric Gk . This formulation makes explicit how PACE differs from standard OT. If Gk (x, t) ≡ I, then Eq. (1) reduces to the Euclidean kinetic energy. For any fixed endpoint pair (x, y), the minimum-action path is the straight constant-speed interpolant, and the induced cost becomes cGk (x, y) = ∥y − x∥2 . In this special case, Eq. (2) reduces to standard quadratic-cost OT between ρ̂k and ρ̂k+1 . PACE departs from this Euclidean case by using a state- and time-dependent anisotropic metric Gk (x, t). The induced cost cGk (x, y) is no longer determined only by endpoint distance; it depends on the entire path and on how the path velocity aligns with local developmental geometry. As a result, the minimum-action cost generally has no closed-form solution and cannot be reduced to a fixed Euclidean endpoint cost. The next section defines the spatiotemporal metric Gk (x, t), and § 3.4 describes how PACE approximates this problem with neural bridges and alternating coupling updates. Finite-dimensional approximation. Problem (2) is an ideal infinite-dimensional formulation. PACE makes it tractable through two approximations (§3.4): (i) the path family is parameterized by neural bridges γθ that satisfy endpoint constraints by construction; (ii) the action integral is approximated on a finite time grid. These yield a finite-dimensional alternating optimization over bridge parameters and correspondence variables. 3.3
Time-dependent spatiotemporal tangent metric
The metric should encode a simple prior: admissible motion should follow directions locally supported by the observed geometry, rather than cut across the cell-state manifold. Although local neighborhoods do not reveal the true direction of time, their dominant variation directions provide an undirected tangent approximation to the set of plausible state changes. PACE therefore treats tangent motion as low cost and penalizes velocity components in the locally estimated normal subspace. Because the local composition and geometry of snapshots can change across experimental time, the admissible subspaces are indexed by both state and time, yielding a time- and state-dependent metric Gk (x, t). Local normal subspaces. At each anchor point xr observed at time tr , PACE estimates an anchorwise local normal direction by finding its mnn nearest spatial neighbors within the same snapshot (including xr itself), computing a Gaussian-kernel weighted covariance of the neighbor cloud, and extracting the minimum-variance principal direction nr . In two dimensions the normal projector is (r) simply PN = nr n⊤ r ; in higher dimensions PACE adaptively selects the tangent-subspace dimension and builds the normal projector as the orthogonal complement (Appendix K). 4
Spatiotemporal metric construction. For a query point (x, t) in segment k, PACE interpolates these anchor-wise local normal projectors across state and time using space-time Gaussian weights over all anchor points r: ! ∥x − xr ∥2 |t − tr |2 (k) ωr (x, t) ∝ exp − − , (3) (k) (k) (hx )2 (ht )2 P (k) (k) (k) with normalization chosen so that r ωr (x, t) = 1, where the bandwidths hx and ht are estimated adaptively for each segment (see Appendix L). These weights define an averaged normal projector field and the corresponding metric tensor X (k) (r) (k) CN (x, t) = ωr(k) (x, t)PN , Gk (x, t) = I + αCN (x, t), α > 0. r (k)
(k)
The corresponding velocity cost is v ⊤ Gk (x, t)v = ∥v∥2 + α v ⊤ CN (x, t) v. Because CN is a positive-semidefinite average of local normal projectors, the second term increases the cost most strongly for velocities aligned with nearby estimated normal directions. Proposition 2 (Normal-subspace penalizing property). For every segment k, query point (x, t), and velocity v ∈ Rd , (k) v ⊤ Gk (x, t)v = ∥v∥2 + αv ⊤ CN (x, t)v ≥ ∥v∥2 . (4) (k)
⊥ Moreover, if CN (x, t) = PN (x, t) is an exact orthogonal projector onto the normal subspace Tx,t , ⊥ and v = vT + vN with vT ∈ Tx,t and vN ∈ Tx,t , then
v ⊤ Gk (x, t)v = ∥vT ∥2 + (1 + α)∥vN ∥2 .
(5)
Proposition 2 shows that the interpolated metric always preserves at least the Euclidean velocity cost, and that tangent motion retains its Euclidean cost while normal motion is penalized by a factor 1 + α in the ideal projector case. Thus, the metric turns the qualitative principle of geometry-aligned cellular motion into an explicit action functional. 3.4
Finite-dimensional optimization by alternating bridge and OT coupling updates
PACE approximates the anisotropic bridge problem by alternating between an endpoint-conditioned bridge and a cross-time coupling. The coupling determines which endpoint pairs are used to train the bridge, while the bridge induces a path-action cost for updating the coupling. This creates a bootstrap mechanism in which more plausible couplings provide better endpoint supervision, and better bridges provide a geometry-aware approximation to the endpoint cost used for OT refinement. Since solving a separate minimum-action path problem for every candidate endpoint pair is infeasible under the state- and time-dependent metric Gk (x, t), PACE amortizes path optimization with a shared endpoint-preserving neural bridge γθ (x, y, τ ). See Appendix J for a detailed discussion of this alternating procedure. Endpoint-preserving neural bridge. For an endpoint pair (x, y) ∈ Rd × Rd and local time τ ∈ [0, 1], we define γθ (x, y, τ ) = (1 − τ )x + τ y + τ (1 − τ )ψθ (x, y, τ ),
(6)
where ψθ : Rd × Rd × [0, 1] → Rd is a neural network. The factor τ (1 − τ ) enforces γθ (x, y, 0) = x and γθ (x, y, 1) = y, so the network only controls the interior deformation of the path [30]. The bridge velocity is uθ (x, y, τ ) = ∂τ γθ (x, y, τ ), computed by automatic differentiation. When ψθ = 0, the bridge reduces to the straight Euclidean interpolant. Bridge learning under fixed coupling. Assume that, for each adjacent snapshot pair, a coupling (k) (k+1) πk ∈ Π(ρ̂k , ρ̂k+1 ) is given. During the bridge update, endpoint pairs (xi , xj ) ∼ πk are sampled from the current coupling. We write tk,τ = tk + τ (tk+1 − tk ). PACE trains the endpoint-preserving bridge by minimizing Lbridge (θ) = λmetric Lmetric + λreg Lreg . 5
(7)
The main term Lmetric is the anisotropic bridge action induced by the spatiotemporal metric Gk . An optional regularization term Lreg penalizes large normal components and incoherent cross-segment velocities (Appendix G). The metric-action loss is " Lmetric = Ek,(x,y)∼πk
# T 1X uθ (x, y, τℓ )⊤ Gk (γθ (x, y, τℓ ), tk,τℓ ) uθ (x, y, τℓ ) . T
(8)
ℓ=1
Coupling update under fixed bridge. After the bridge model has been updated under the current coupling, PACE recomputes the coupling using the full path action rather than Euclidean endpoint (k) (k+1) distance. For a candidate source-target pair (xi , xj ), the path-action cost is M 1 X (k) (k+1) (k) (k+1) (k) (k+1) uθ (xi , xj , τm )⊤ Gk γθ (xi , xj , τm ), tk,τm uθ (xi , xj , τm ). M m=1 (9) The coupling is then updated by solving X path (k) πk ← arg min cij (θ)πij . (10)
cpath ij (θ) =
πk ∈Π(ρ̂k ,ρ̂k+1 )
Alternating optimization.
i,j
Given an initial coupling, PACE alternates between two blocks: (k)
(k+1)
1. Bridge update. Fix the current couplings {πk }, sample endpoint pairs (xi , xj and update θ by taking gradient steps on Eq. (7).
) ∼ πk ,
2. Coupling update. Fix the current bridge γθ , compute the path-action cost matrix in Eq. (9), and update each πk by solving the OT problem in Eq. (10). The bridge update learns low-action endpoint-conditioned paths under the current coupling, while the coupling update selects source-target transport using the learned path action. In the idealized finite-dimensional setting where the bridge block is solved exactly for the metric-action objective, these two updates form a monotone block-coordinate descent scheme; Appendix F states this property and relates it to the stochastic neural implementation. In practice, the coupling update is triggered periodically, for example every R0 epochs. 3.5
Distilling endpoint-conditioned bridges into a global velocity field
The alternating optimization in §3.4 yields endpoint-conditioned bridges γθ (x, y, τ ) and τ -velocities uθ (x, y, τ ) = ∂τ γθ (x, y, τ ). These bridges are pair-specific: to query a velocity at an arbitrary state z and time t, one would need to know which endpoint pair (x, y) and interpolation parameter τ generated z. PACE therefore distills the bridge dynamics into a global velocity field vϕ (x, t) that can be evaluated without reference to a particular endpoint pair. Distillation objective. For each segment k, we sample (x, y) ∼ πk and τ ∼ Uniform[0, 1], set xtk,τ = γθ (x, y, τ ) with tk,τ = tk + τ ∆tk , and match the global velocity field: " # 2 uθ (x, y, τ ) Ldistill (ϕ) = Ek,(x,y)∼πk Eτ vϕ (xtk,τ , tk,τ ) − , (11) ∆tk where ∆tk = tk+1 − tk converts τ -velocity to physical-time velocity. Once vϕ is learned, continuous trajectories are generated by integrating dxt = vϕ (xt , t), dt
6
x0 ∼ ρ̂0 .
(12)
Oceans
5
24
4
4
15
20
3
3
2
2
1
1
16
10
12 8
5
4
0
0
0
0
EB(PHATE)
Two − Branch
Schiebinger2019
Time
8 7 6 5 4 3 2 1 0
iPSC − Liu(PCA)
Figure 2: Overview of the 2D benchmark datasets. Points are colored by observed time for Ocean [30], Two-Branch, EB PHATE [18], Schiebinger2019 [3], and iPSC-Liu [21]. Table 1: Per-timepoint results on Ocean (2D) [30], holdout t ∈ {1, 3, 5, 7}. MMD ↓
Method
W1 ↓
W2 ↓
t=1
t=3
t=5
t=7
t=1
t=3
t=5
t=7
t=1
t=3
t=5
t=7
Action Matching Aligned CFM CURLY DMSB MFM OT-CFM
0.7429 0.7989 1.0545 1.1260 0.8482 0.9410
1.0662 0.8754 0.8373 0.6089 0.7859 0.8169
0.9465 1.0073 0.5869 0.2542 0.6928 0.6827
1.1289 0.8272 0.8663 0.4962 0.5823 0.6305
0.3549 0.0976 0.3766 0.3060 0.0960 0.1227
0.4823 0.1644 0.1521 0.1097 0.1433 0.1428
0.2383 0.3024 0.1077 0.0435 0.1359 0.1289
0.5073 0.2245 0.2638 0.1097 0.1279 0.1321
0.3662 0.1017 0.3771 0.3065 0.0991 0.1255
0.4846 0.1673 0.1537 0.1150 0.1452 0.1459
0.2425 0.3046 0.1102 0.0473 0.1385 0.1308
0.5143 0.2262 0.2646 0.1142 0.1302 0.1334
PACE(ours)
0.4504 0.2588 0.0735 0.2779 0.0399 0.0398 0.0268 0.0505 0.0440 0.0434 0.0362 0.0534
4
Experiment
The experiments evaluate PACE through five questions. First, on controlled 2D trajectories, can PACE recover smooth and branching geometry without paired identities or velocity supervision (Figure 2; Tables 1 and A.2)? Second, when reference velocities are available only for evaluation, do the learned dynamics align with held-out velocity fields on Ocean and EB PHATE (Figure 3; Appendix Tables A.3 and A.4)? Third, on biological time courses, does PACE improve held-out reconstruction for single-cell differentiation and reprogramming from destructive snapshots (Table 2)? Fourth, does PACE remain effective as the representation dimension increases (Tables 3)? Finally, which components of PACE account for the observed gains (Figure 5)? 4.1
Experimental setup
Evaluation protocol. All experiments use the same held-out snapshot protocol: methods observe only the training time points and are evaluated on the held-out time points listed in each table caption. We report distributional reconstruction quality using maximum mean discrepancy (MMD), 1-Wasserstein distance (W1 ), and 2-Wasserstein distance (W2 ); lower is better, bold denotes the best result, and underline denotes the second best. Datasets. Our benchmark suite includes controlled 2D temporal point clouds and biological single-cell time-course datasets. The controlled datasets are designed to isolate geometric behavior in low-dimensional temporal data, including Ocean [35, 30] and a branching toy dataset, without relying on cell identities or velocity supervision. The biological datasets test destructive snapshot settings in which individual cell identities are not shared across time, covering embryoid body differentiation [18], mouse reprogramming [3], induced trophoblast stem-cell reprogramming [21], and multimodal CITE-seq/Multiome time courses [36]. Figure 2 gives a visual overview of the 2D datasets used in the main low-dimensional experiments. Detailed dataset sources and preprocessing choices are provided in Appendix A. Baselines. We compare against representative transport- and flow-based trajectory models, covering action-based continuous dynamics [14], minibatch optimal-transport flow matching [15], metric/geodesic flow matching [34], pseudo-velocity-corrected flows [30], adversarial multi-marginal interpolant learning [37], and stochastic Schrödinger bridge dynamics [29]. Implementation details for baselines and PACE are provided in Appendices B and C. Qualitative prediction-versus-held-out overlays for the same 2D runs are shown in Appendix Figures A.1–A.5. 7
Controlled 2D trajectories Cosine Distance
0.20
Normalized L2
0.6
Cosine Distance
Normalized L2
Cosine Distance
0.5 0.15
0.4 0.3
0.10
0.2
0.05
0.1
0.00
n Actio
hing CFM CURLY Matc Aligned
FM
OT-C
MFM
urs)
(O PACE
0.0
(a) Ocean
0.8 0.7 0.6 0.5 0.4 0.3 0.2 0.1 0.0
Cosine Distance
Normalized L2
1.2 1.0 0.8
Normalized L2
4.2
0.6 0.4 0.2
MFM
Y CURL
FM
OT-C
g rs) FM ed C n MatchinPACE(Ou Actio
0.0
Align
(b) EB PHATE
Figure 3: Velocity-alignment diagnostics on Ocean [30] and EB PHATE [18]. Ocean compares learned velocities with held-out simulator velocities, while EB PHATE compares learned 2D velocities with the RNA-velocity reference. Bars report cosine distance and normalized L2 error; lower is better. We evaluate PACE on two controlled 2D settings, including Ocean, a rotational non-gradient benchmark from Curly Flow Matching [30], and a Two-Branch toy benchmark. Methods observe only unpaired point positions; simulator velocities and latent identities are not used for training. For Ocean, we also hold out evaluation snapshots and compare the learned velocity field at those times with the simulator velocity field, testing whether the inferred dynamics recover the correct directionality rather than only matching held-out marginals. Tables 1 and A.2 show that PACE gives the strongest overall held-out reconstruction. It is best on all Ocean metrics and remains strongest on most Two-Branch entries, with CURLY slightly better only for late-branch W2 . The Ocean panel in Figure 3 further checks the learned velocity direction against held-out simulator velocities; the per-timepoint diagnostic values are reported in Appendix Table A.3. Together, these results indicate that the geometry-aware bias helps recover both smooth rotational motion and simple branching structure from snapshots alone. 4.3
Low-dimensional single-cell trajectories
We next ask whether the same behavior carries over to biological time courses represented in lowdimensional embeddings, where the true cell identities are unobserved and geometric structure must be inferred from destructive snapshots. Table 2 summarizes the time-averaged low-dimensional biological results across EB PHATE, iPSC-Liu, and Schiebinger2019; the corresponding per-timepoint results are reported in Appendix Tables A.5, A.6, and A.7. On EB PHATE, PACE obtains the best MMD, W1 , and W2 , showing that the method improves held-out marginal reconstruction in a manifold-aware PHATE embedding. The EB PHATE panel in Figure 3 adds an independent directionality check by comparing the learned 2D velocity field at the held-out snapshot with RNA velocity, which is used only for evaluation. Appendix Table A.4 shows that PACE also best aligns with this external velocity proxy, suggesting that the gain reflects both endpoint reconstruction and plausible local flow direction. On iPSC-Liu 2D, PACE is best on every reported metric. On Schiebinger2019, PACE is best on seven of nine detailed metric-timepoint entries and remains close on the remaining final-time scores. Overall, the low-dimensional biological results show that PACE transfers the controlled-geometry gains to both differentiation and reprogramming snapshots. Compared with the strongest baseline for each dataset–metric pair, PACE reduces MMD, W1 , and W2 by 36.2%, 21.5%, and 13.6% on EB PHATE; by 16.1%, 26.4%, and 25.0% on iPSC-Liu; and by 31.2%, 25.5%, and 22.4% on Schiebinger2019.
8
Table 2: Time-averaged low-dimensional single-cell results. EB [18] uses t = 3, iPSC-Liu [21] uses t ∈ {4, 16}, and Schiebinger2019 [3] uses t ∈ {6, 11, 16}. Method
EB PHATE
iPSC-Liu
Schiebinger2019
MMD ↓ W1 ↓
W2 ↓ MMD ↓ W1 ↓
W2 ↓ MMD ↓ W1 ↓
W2 ↓
Action Matching Aligned CFM CURLY DMSB MFM OT-CFM
0.2436 0.1056 0.1384 0.1325 0.1518 0.1069
0.4553 0.2657 0.3223 0.3053 0.3704 0.2784
0.5273 0.3243 0.4422 0.4341 0.4451 0.3622
0.5230 0.5428 0.5613 0.6456 0.5533 0.5228
0.8005 0.6594 0.7186 1.0244 0.6657 0.7166
0.9730 0.8320 0.9216 1.1023 0.8838 0.9261
0.4057 0.3239 0.7687 0.3030 0.6990 0.8072
0.4161 0.3175 1.3684 0.2774 1.5897 1.9095
0.4754 0.3931 1.4660 0.3662 1.7942 2.0235
PACE(ours)
0.0674
0.2087 0.2803
0.4385
0.4853 0.6238
0.2084
0.2067 0.2840
4.4
Higher-dimensional single-cell benchmarks
High-dimensional PCA tests whether PACE still helps Table 3: iPSC-Liu [21] time-averaged results when Euclidean distances lose temporal contrast. over holdouts t ∈ {4, 8, 12, 16}. On iPSC-Liu, PACE is best on all time-averaged Dim Method MMD ↓ W1 ↓ W2 ↓ 10D/50D metrics in Table 3. On OP-Cite/OP-Multi 100D, it remains on the empirical Pareto front, leading 10D Action Matching 0.5555 3.7801 4.3743 Aligned CFM 0.4644 2.9405 3.3825 OP-Cite MMD and OP-Multi W1 /W2 in Appendix CURLY 0.5369 3.6128 4.2102 Table A.9. Figure 4 shows the diagnostic behind this DMSB 0.7713 6.3088 6.4593 regime, with norm CV approaching 0.3 and inter-time MFM 0.5273 3.4066 3.8382 OT-CFM 0.5145 3.3850 3.9447 separation approaching the within-time radius. The PACE (ours) 0.4095 2.9226 3.1879 useful signal is therefore not global separation between time means, but local anisotropy within each 50D Action Matching 0.3125 7.6126 8.4868 Aligned CFM 0.2624 7.3133 8.0552 snapshot; PACE turns that local geometry into couCURLY 0.3144 8.1144 8.9096 pling costs. This helps distinguish directionally conDMSB 0.5638 10.9368 11.2423 sistent moves from distance-similar but geometrically MFM 0.2853 7.7924 8.5454 OT-CFM 0.2898 7.8392 8.5954 implausible pairings. Appendix H and Figures A.6– PACE (ours) 0.2550 6.7537 7.2836 A.7 provide the detailed concentration analysis. 4.5
Ablation study
We finally isolate the contribution of the main PACE components on the Schiebinger2019 holdouts used in Table A.7. Figure 5 compares the full model with variants that remove coupling rematching, set the metric penalty to the Euclidean case, use all neighbors rather than local neighborhoods, or ablate the spatial and temporal kernel structure used to smooth local geometry. Each simplification degrades at least one metric or time point, and the full PACE variant gives the most stable low errors across MMD, W1 , and W2 . This indicates that the gains do not come from a single implementation detail: the anisotropic metric, local-neighborhood construction, spatial-temporal geometry smoothing, and rematching step all contribute to the final trajectory reconstruction.
median kci −cj k / r̄ij
0.5
101
PCA dimension d
102
0.35
MMD
0.4
0.30
0.3 0.2
PCA dimension d
102
0.25 0.20
0.1 101
Spatial Kernel Temporal Kernel
W2
1.8 1.6 1.4 1.2 1.0 0.8 0.6
CV(kxk) = Std(kxk)/ [kxk]
0.55 0.50 0.45 0.40 0.35 0.30 0.25
=0 (Euclidean) KNN = All
PACE No Rematch
iPSC warning threshold
W1
op-cite op-multi
6
11 Test Timepoint
16
0.15
6
11 Test Timepoint
16
0.45 0.40 0.35 0.30 0.25 0.20
6
11 Test Timepoint
16
Figure 4: High-dimensional concentra- Figure 5: PACE ablations on Schiebinger2019 [3] holdouts. tion diagnostics. Dashed lines mark Lower is better; variants remove rematching, use Euclidean norm CV = 0.3 and inter-time/within- action (α = 0), use all neighbors, or drop spatial/temporal time ratio = 1.0. kernels. 9
5
Conclusion
In this paper, we have introduced PACE, a geometry-aware framework for trajectory inference from destructive single-cell time-course snapshots. PACE replaces fixed Euclidean endpoint costs with an anisotropic path cost induced by a local spatiotemporal metric, then alternates bridge learning with coupling refinement and distills the resulting dynamics into a velocity field. Across controlled and biological benchmarks, the results support local geometry as a useful inductive bias for held-out snapshot reconstruction without paired cells, lineage tracing, or velocity supervision; the velocity diagnostics further show that the learned flows recover plausible directionality when external velocity references are available only for evaluation. Ablations show that this behavior depends on the combination of anisotropic local costs, neighborhood-restricted geometry, smoothing, and coupling rematching rather than a single implementation detail. Future work will focus on more robust metric estimation, uncertainty quantification, and partially supervised or multimodal time-course settings.
Acknowledgement We thank Yuchen Yang, Minsi Ren, Zhuo Xu, Yucheng Luo, and Mingzheng Fang for discussions and for providing feedback on our manuscript. We also gratefully acknowledge the support of Westlake University Research Center for Industries of the Future; Westlake University Center for High-performance Computing. The content is solely the responsibility of the authors and does not necessarily represent the official views of the funding entities.
References [1] Jeffrey A Farrell, Yiqun Wang, Samantha J Riesenfeld, Karthik Shekhar, Aviv Regev, and Alexander F Schier. Single-cell reconstruction of developmental trajectories during zebrafish embryogenesis. Science, 360(6392):eaar3131, 2018. [2] Daniel E Wagner, Caleb Weinreb, Zach M Collins, James A Briggs, Sean G Megason, and Allon M Klein. Single-cell mapping of gene expression landscapes and lineage in the zebrafish embryo. Science, 360(6392):981–987, 2018. [3] 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. [4] Amos Tanay and Aviv Regev. Scaling single-cell genomics from phenomenology to mechanism. Nature, 541(7637):331–338, 2017. [5] Jonathan A Griffiths, Antonio Scialdone, and John C Marioni. Using single-cell genomics to understand developmental processes and cell fate decisions. Molecular systems biology, 14(4):MSB178046, 2018. [6] Cole Trapnell, Davide Cacchiarelli, Jonna Grimsby, Prapti Pokharel, Shuqiang Li, Michael Morse, Niall J Lennon, Kenneth J Livak, Tarjei S Mikkelsen, and John L Rinn. The dynamics and regulators of cell fate decisions are revealed by pseudotemporal ordering of single cells. Nature biotechnology, 32(4):381–386, 2014. [7] Laleh Haghverdi, Maren Büttner, F Alexander Wolf, Florian Buettner, and Fabian J Theis. Diffusion pseudotime robustly reconstructs lineage branching. Nature methods, 13(10):845– 848, 2016. [8] Manu Setty, Vaidotas Kiseliovas, Jacob Levine, Adam Gayoso, Linas Mazutis, and Dana Pe’Er. Characterization of cell fate probabilities in single-cell data with palantir. Nature biotechnology, 37(4):451–460, 2019. [9] Davide Cacchiarelli, Xiaojie Qiu, Sanjay Srivatsan, Anna Manfredi, Michael Ziller, Eliah Overbey, Antonio Grimaldi, Jonna Grimsby, Prapti Pokharel, Kenneth J Livak, et al. Aligning single-cell developmental and reprogramming trajectories identifies molecular determinants of myogenic reprogramming outcome. Cell Systems, 7(3):258–268, 2018. 10
[10] Hugo Lavenant, Stephen Zhang, Young-Heon Kim, and Geoffrey Schiebinger. Towards a mathematical theory of trajectory inference. arXiv preprint arXiv:2102.09204, 2021. [11] Tatsunori Hashimoto, David Gifford, and Tommi Jaakkola. Learning population-level diffusions with generative rnns. In International Conference on Machine Learning, pages 2417–2426. PMLR, 2016. [12] Gabriel Peyré and Marco Cuturi. Computational optimal transport: With applications to data science. Now Foundations and Trends, 2019. [13] Caleb Weinreb, Samuel Wolock, Betsabeh K Tusi, Merav Socolovsky, and Allon M Klein. Fundamental limits on dynamic inference from single-cell snapshots. Proceedings of the National Academy of Sciences, 115(10):E2467–E2476, 2018. [14] 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. [15] Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian Fatras, Guy Wolf, and Yoshua Bengio. Conditional flow matching: Simulation-free dynamic optimal transport. arXiv preprint arXiv:2302.00482, 2(3), 2023. [16] 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. [17] Guillaume Huguet, Daniel Sumner Magruder, Alexander Tong, Oluwadamilola Fasina, Manik Kuchroo, Guy Wolf, and Smita Krishnaswamy. Manifold interpolating optimal-transport flows for trajectory inference. Advances in neural information processing systems, 35:29705–29718, 2022. [18] 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. [19] F Alexander Wolf, Fiona K Hamey, Mireya Plass, Jordi Solana, Joakim S Dahlin, Berthold Göttgens, Nikolaus Rajewsky, Lukas Simon, and Fabian J Theis. Paga: graph abstraction reconciles clustering with trajectory inference through a topology preserving map of single cells. Genome biology, 20(1):59, 2019. [20] Sophie Tritschler, Maren Büttner, David S Fischer, Marius Lange, Volker Bergen, Heiko Lickert, and Fabian J Theis. Concepts and limitations for learning developmental trajectories from single cell genomics. Development, 146(12):dev170506, 2019. [21] Xiaodong Liu, John F Ouyang, Fernando J Rossello, Jia Ping Tan, Kathryn C Davidson, Daniela S Valdes, Jan Schroeder, Yu BY Sun, Joseph Chen, Anja S Knaupp, et al. Reprogramming roadmap reveals route to human induced trophoblast stem cells. Nature, 586(7827):101– 107, 2020. [22] Gioele La Manno, Ruslan Soldatov, Amit Zeisel, Emelie Braun, Hannah Hochgerner, Viktor Petukhov, Katja Lidschreiber, Maria E Kastriti, Peter Lönnerberg, Alessandro Furlan, et al. Rna velocity of single cells. Nature, 560(7719):494–498, 2018. [23] Volker Bergen, Marius Lange, Stefan Peidli, F Alexander Wolf, and Fabian J Theis. Generalizing rna velocity to transient cell states through dynamical modeling. Nature biotechnology, 38(12):1408–1414, 2020. [24] Marius Lange, Volker Bergen, Michal Klein, Manu Setty, Bernhard Reuter, Mostafa Bakhti, Heiko Lickert, Meshal Ansari, Janine Schniering, Herbert B Schiller, et al. Cellrank for directed single-cell fate mapping. Nature methods, 19(2):159–170, 2022. [25] Wouter Saelens, Robrecht Cannoodt, Helena Todorov, and Yvan Saeys. A comparison of single-cell trajectory inference methods. Nature biotechnology, 37(5):547–554, 2019. 11
[26] 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. [27] Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in neural information processing systems, 34:17695–17709, 2021. [28] Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching. Advances in neural information processing systems, 36:62183–62223, 2023. [29] Tianrong Chen, Guan-Horng Liu, Molei Tao, and Evangelos Theodorou. Deep momentum multi-marginal schrödinger bridge. Advances in Neural Information Processing Systems, 36:57058–57086, 2023. [30] Katarina Petrović, Lazar Atanackovic, Viggo Moro, Kacper Kapuśniak, Ismail Ilkan Ceylan, Michael M. Bronstein, Joey Bose, and Alexander Tong. Curly flow matching for learning non-gradient field dynamics. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2026. [31] Søren Hauberg, Oren Freifeld, and Michael Black. A geometric take on metric learning. Advances in Neural Information Processing Systems, 25, 2012. [32] Georgios Arvanitidis, Lars Kai Hansen, and Søren Hauberg. Latent space oddity: on the curvature of deep generative models. arXiv preprint arXiv:1710.11379, 2017. [33] Ricky T. Q. Chen and Yaron Lipman. Flow matching on general geometries. In The Twelfth International Conference on Learning Representations, 2024. [34] Kacper Kapuśniak, Peter Potaptchik, Teodora Reu, Leo Zhang, Alexander Tong, Michael Bronstein, Avishek J 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. [35] Yunyi Shen, Renato Berlinghieri, and Tamara Broderick. Multi-marginal schrödinger bridges with iterative reference refinement. In The 28th International Conference on Artificial Intelligence and Statistics, 2025. [36] 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. [37] Oskar Kviman, Kirill Tamogashev, Nicola Branchini, Víctor Elvira, Jens Lagergren, and Nikolay Malkin. Multi-marginal flow matching with adversarially learnt interpolants. In The Fourteenth International Conference on Learning Representations, 2026. [38] Eric P Chassignet, Harley E Hurlburt, Ole Martin Smedstad, George R Halliwell, Patrick J Hogan, Alan J Wallcraft, Remy Baraille, and Rainer Bleck. The hycom (hybrid coordinate ocean model) data assimilative system. Journal of Marine Systems, 65(1-4):60–83, 2007. [39] Michel Ledoux. The concentration of measure phenomenon. Number 89. American Mathematical Soc., 2001. [40] Roman Vershynin. High-dimensional probability: An introduction with applications in data science, volume 47. Cambridge university press, 2018. [41] Kevin Beyer, Jonathan Goldstein, Raghu Ramakrishnan, and Uri Shaft. When is “nearest neighbor” meaningful? In International conference on database theory, pages 217–235. Springer, 1999.
12
Technical appendices and supplementary material A. Benchmark and Metric Details . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 13 B. Baseline Details . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 23 C. PACE Implementation Details . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 24 D. Proof of Proposition 1 (Ill-posedness) . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 24 E. Local Metric and Correspondence . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 24 F. Monotonicity of Ideal Alternating Updates . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 25 G. Stage 1 Regularizers . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 26 H. High-Dimensional Concentration Effects . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .26 I. iPSC-Liu Results . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 28 J. Intuition behind alternating bridge and coupling optimization . . . . . . . . . . . . . . . . . . . . . . . 29 K. Adaptive tangent and normal projectors . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 29 L. Adaptive bandwidth estimation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 30 M. Limitations and Future Work . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 31
A
Benchmark and Metric Details
All datasets are cast as unpaired time-course point clouds. A point is a cell or synthetic particle state, and a time point is an empirical marginal distribution. Methods receive only the training snapshots; held-out snapshots are used only for evaluation. No method receives cell identities, cross-time correspondences, RNA velocity, or any ground-truth trajectory pairing. Shared split and rollout protocol. For a held-out label th , the evaluator finds the nearest observed training labels ta < th < tb . The learned velocity field is rolled out from the empirical source cloud at ta toward tb on the same normalized training-time scale used during training. We use 101 Euler ODE steps and extract the intermediate frame at ratio (th − ta )/(tb − ta ). The predicted cloud is then compared with the empirical held-out cloud at th . Thus, the reported numbers measure recovery of the missing marginal distribution, not recovery of unobserved cell identities.
13
Table A.1: Dataset sources, representations, and held-out splits used in Section 4. Controlled 2D coordinates are used directly. Biological embeddings are standardized by fitting a StandardScaler on training snapshots only, then applying it to train and held-out snapshots. Benchmark
Source and role
Representation and prepro- Train/test labels cessing
Ocean (Gulf of HYCOM-derived particle bench- data/oceans/oceans. Train {0, 2, 4, 6, 8}; Mexico vortex) mark from Shen et al. [35], also npz; positions has shape test {1, 3, 5, 7}. used in CURLY [30]. Training (9, 111, 2). No whitening; uses only particle positions; veloci- full frames are used. ties are not used for supervision. Two-Branch
In-house synthetic bifurca- Six 2D frames, 64 points Train {0, 2, 3, 5}; test tion generated by src/data_ per frame, with latent pseu- {1, 4}. preprocess/generate_two_ dotime and branch labels branch_data.py. It tests recov- stored for diagnostics but ery of a trunk that splits into two not given to methods. No branches. whitening; full frames are used.
EB PHATE
Embryoid body differentiation data data/eb_velocity_v5. Train {0, 1, 2, 4}; test from PHATE [18]. It represents npz; 16,819 cells with {3}. a low-dimensional differentiation phate (N, 2) and pcs trajectory. (N, 100). The 2D experiment uses PHATE and fits whitening on training cells only.
Schiebinger2019 Mouse reprogramming time course Local H5AD file; X stores Train all integer days from Schiebinger et al. [3], con- 2D FLE coordinates and 0, . . . , 18 except verted from the scNODE data pack- obs["day"] stores integer {6, 11, 16}; test age using serum-filtered FLE coor- day labels. Whitening is fit {6, 11, 16}. dinates. on training cells only. iPSC-Liu
Human reprogramming time data/iPSC_liu_pca. 2D: train course from Liu et al. [21]. It is h5ad; obsm["X_pca"] {0, 8, 12, 20, 24}, test used to evaluate missing-timepoint provides PCA scores and {4, 16}. 10D/50D: reconstruction in PCA representa- obs["timepoint"] stores separate singletions. labels such as D0, parsed as holdout runs for integers. Whitening is fit on t ∈ {4, 8, 12, 16}. training cells only.
OP-Cite/OPMulti
NeurIPS Open Problems multi- op_cite_inputs_0.h5ad Train {2, 4, 7}; test modal single-cell integration chal- and op_train_multi_ {3}. lenge data [36]. OP-Cite uses targets_0.h5ad; both CITE-seq features; OP-Multi uses use obsm["X_pca"]. Multiome features. Table A.9 uses all 100 PCs, with whitening fit on training cells only.
Ocean benchmark. The Ocean benchmark follows the Gulf of Mexico vortex experiment of Shen et al. [35]. The data are generated by first extracting a velocity field around a vortex feature from high-resolution HYbrid Coordinate Ocean Model (HYCOM) reanalysis data [38], then simulating particles representing buoys or ocean debris through that field. The original benchmark contains approximately 1000 observations across five training times and four validation times; the released array used here has nine snapshots with 111 particles per snapshot. Unlike the reference-family model in Shen et al., PACE and all baselines in our comparison receive only the particle positions at training snapshots, with no velocity supervision. The HYCOM-derived velocity field is held out from training and used only for the velocity-alignment diagnostic in Table A.3. Velocity-alignment diagnostics. Figure 3 summarizes the velocity-alignment checks used in the main experiments. The tables below report the numerical values behind the two panels: Ocean compares learned velocities with held-out simulator velocities, while EB PHATE compares learned 2D 14
velocities with the RNA-velocity reference. These reference velocities are used only for diagnostics, not for training. Table A.2: Branching toy (2D) results at held-out time points t ∈ {1, 4}. MMD ↓ t=1 t=4
Method
W1 ↓ t=1 t=4
W2 ↓ t=1 t=4
Action Matching 0.6129 0.2905 0.3343 0.2732 0.3419 0.3048 Aligned CFM 0.6524 0.3439 0.3874 0.309 0.3944 0.3371 CURLY 0.2664 0.0767 0.1703 0.1797 0.2123 0.2087 DMSB 1.0232 1.1585 4.5832 2.8358 4.5845 2.8406 MFM 0.3817 0.1145 0.1686 0.1865 0.1832 0.2116 OT-CFM 0.5279 0.1244 0.2433 0.1841 0.2509 0.2119 PACE(ours)
0.1968 0.0146 0.1386 0.1752 0.1686 0.2236
Table A.3: Ocean [30] velocity-alignment diagnostics on held-out time points. The simulator velocity is used only for evaluation. Cosine distance is 1 − cos(v̂, v), and normalized L2 compares unitnormalized velocity directions; lower is better. Cosine distance ↓
Method
Normalized L2 ↓
t=1 t=3 t=5 t=7
Avg.
t=1 t=3 t=5 t=7
Avg.
Action Matching Aligned CFM CURLY MFM OT-CFM
0.3535 0.0440 0.0227 0.0029 0.0040
0.1954 0.0221 0.0186 0.0056 0.0072
0.8174 0.2808 0.1997 0.0417 0.0738
0.5612 0.1840 0.1642 0.0748 0.0854
PACE(ours)
0.0041 0.0030 0.0069 0.0031 0.0043 0.0756 0.0631 0.0671 0.0633 0.0673
0.0408 0.0269 0.0068 0.0063 0.0101
0.1974 0.0101 0.0220 0.0089 0.0088
0.1900 0.0073 0.0231 0.0043 0.0058
0.2317 0.2118 0.0906 0.0913 0.1147
0.5965 0.1311 0.1709 0.0965 0.0731
0.5993 0.1122 0.1955 0.0698 0.0801
Table A.4: EB PHATE [18] velocity-alignment diagnostics at the held-out snapshot t = 3. The RNA-velocity reference is used only for evaluation. Cosine distance is 1 − cos(v̂, v), and normalized L2 compares unit-normalized velocity directions; lower is better.
Method
Cosine distance ↓ Normalized L2 ↓
Action Matching Aligned CFM CURLY MFM OT-CFM
0.4041 0.4273 0.6325 0.6930 0.5571
0.7542 0.7707 0.9693 1.0261 0.8807
PACE(ours)
0.3287
0.6630
Low-dimensional single-cell reconstruction results. Table 2 reports the time-averaged summary used in the main text. The tables below give the corresponding per-timepoint results for each lowdimensional biological benchmark: EB PHATE has one held-out PHATE snapshot, iPSC-Liu has two held-out 2D PCA snapshots, and Schiebinger2019 has three held-out reprogramming days. These detailed tables show whether the averaged gains are consistent across individual held-out times. Subsampling and dimensionality. For biological datasets, samples_per_timepoint is either an experiment-specific cap or null. When it is a cap, each selected time point is sampled without replacement using the experiment seed before whitening. When it is null, all available cells from the selected time points are used. For EB, the raw timepoint counts are 2381, 4163, 3278, 3665, 3332 across labels 0, . . . , 4. For iPSC-Liu and OP-Cite/OP-Multi, the dimensionality is chosen by taking the first d PCA coordinates from X_pca. The 10D and 50D iPSC-Liu tables aggregate separate single-holdout runs; in each run the held-out label is removed from the training label set. Metrics. For every held-out time point, we compare the predicted point cloud with the observed point cloud using three distributional metrics. MMD is computed with a Gaussian kernel and a 15
median-heuristic bandwidth, and the reported value is the square root of the nonnegative MMD estimate. W1 uses uniform empirical weights and the Euclidean ground cost between predicted and observed samples. W2 uses uniform empirical weights and squared Euclidean ground cost, then reports the square root of the optimal transport value. These metrics are distributional: they evaluate reconstruction of the held-out marginal distribution, not cell-wise recovery of unobserved identities. Table A.5: Per-timepoint results on EB PHATE [18] (2D) with held-out t = 3. MMD ↓ W1 ↓
W2 ↓
t=3
t=3
t=3
Action Matching Aligned CFM CURLY DMSB MFM OT-CFM
0.2436 0.1056 0.1384 0.1325 0.1518 0.1069
0.4553 0.5273 0.2657 0.3243 0.3223 0.4422 0.3053 0.4341 0.3704 0.4451 0.2784 0.3622
PACE(ours)
0.0674 0.2087 0.2803
Method
Table A.6: Per-timepoint results on iPSC-Liu [21] (2D) with held-out time points t ∈ {4, 16}. MMD ↓
Method
W1 ↓
W2 ↓
t=4
t = 16
t=4
t = 16
t=4
t = 16
Action Matching Aligned CFM CURLY DMSB MFM OT-CFM
0.7369 0.8265 0.8156 0.8663 0.8334 0.7787
0.3090 0.2590 0.3069 0.4248 0.2731 0.2668
0.9719 0.7781 0.9030 0.7176 0.7646 0.8695
0.6290 0.5407 0.5341 1.3311 0.5667 0.5636
1.0726 0.7858 0.9292 0.7986 0.7860 0.8983
0.8733 0.8781 0.9140 1.4060 0.9815 0.9539
PACE(ours)
0.6808 0.1962 0.5579 0.4126 0.5646 0.6830
Table A.7: Per-timepoint results on Schiebinger2019 [3], holdout t ∈ {6, 11, 16}. Method MMD ↓ W1 ↓ W2 ↓ t=6
t = 11 t = 16
t=6
t = 11 t = 16
t=6
t = 11 t = 16
Action Matching Aligned CFM CURLY DMSB MFM OT-CFM
0.5056 0.5712 0.9308 0.5081 0.8676 1.1550
0.4940 0.3405 1.0201 0.3196 0.8858 1.0005
0.3203 0.4126 1.4737 0.3118 2.6083 3.1652
0.4982 0.3173 1.9135 0.3108 1.4602 1.8933
0.3283 0.4291 1.5113 0.3174 3.0265 3.2608
0.5337 0.3737 1.9386 0.3760 1.4793 1.9239
PACE(ours)
0.4219 0.1373 0.0660 0.2609 0.1590 0.2003 0.2669 0.2037 0.3813
0.2174 0.0601 0.3553 0.0813 0.3437 0.2661
16
0.4298 0.2228 0.7182 0.2096 0.7005 0.6701
0.5642 0.3764 0.9481 0.4051 0.8768 0.8856
Qualitative Stage-2 trajectory visualizations. Figures A.1–A.5 complement the quantitative tables with selected 2D Stage-2 rollouts. For each dataset, the trajectory panel shows the learned rollout geometry for each baseline, while the prediction panel overlays the predicted held-out marginal with the empirical held-out snapshot. These plots are qualitative diagnostics only; all models are trained from unpaired snapshots without cell identities or ground-truth trajectories.
17
Action Matching
Aligned CFM
CURLY
MFM
OT-CFM
PACE
DMSB
Train anchors Stage 2 trajectory Held-out true Prediction
Action Matching
Aligned CFM
CURLY
MFM
OT-CFM
PACE
DMSB
Held-out true Prediction
Figure A.1: Qualitative Stage-2 results on Ocean [35, 30]. The top panel shows baseline rollouts; the bottom panel compares predicted and held-out point clouds.
18
Action Matching
Aligned CFM
CURLY
MFM
OT-CFM
PACE
DMSB
Train anchors Stage 2 trajectory Held-out true Prediction
Action Matching
Aligned CFM
CURLY
MFM
OT-CFM
PACE
DMSB
Held-out true Prediction
Figure A.2: Qualitative Stage-2 results on the Two-Branch toy benchmark. The top panel shows baseline rollouts; the bottom panel compares predicted and held-out point clouds.
19
Action Matching
Aligned CFM
CURLY
MFM
OT-CFM
PACE
DMSB
Train t=0 Train t=1 Train t=2 Train t=4 Stage 2 trajectory Held-out true t=3 Prediction t=3
Action Matching
Aligned CFM
CURLY
MFM
OT-CFM
PACE
DMSB
Held-out true t=3 Prediction t=3
Figure A.3: Qualitative Stage-2 results on EB PHATE [18]. The top panel shows baseline rollouts; the bottom panel compares predicted and held-out point clouds.
20
Action Matching
Aligned CFM
CURLY
MFM
OT-CFM
PACE
DMSB
Train anchors Stage 2 trajectory Held-out true Prediction
Action Matching
Aligned CFM
CURLY
MFM
OT-CFM
PACE
DMSB
Held-out true Prediction
Figure A.4: Qualitative Stage-2 results on Schiebinger2019 [3]. The top panel shows baseline rollouts; the bottom panel compares predicted and held-out point clouds.
21
Action Matching
Aligned CFM
CURLY
MFM
OT-CFM
PACE
DMSB
Train anchors Stage 2 trajectory Held-out true Prediction
Action Matching
Aligned CFM
CURLY
MFM
OT-CFM
PACE
DMSB
Held-out true Prediction
Figure A.5: Qualitative Stage-2 results on iPSC-Liu [21]. The top panel shows baseline rollouts; the bottom panel compares predicted and held-out point clouds.
22
B
Baseline Details
The result tables include only methods that were run for the corresponding experiment. All deterministic flow baselines use the same dataloaders, train/test labels, whitening rules, and held-out ODE rollout evaluator described in Appendix A. Unless an experiment-specific config overrides it, flow networks are MLP velocity fields trained with Adam/AdamW-style optimizers. The formulas below describe the implemented objectives under a segment [tk , tk+1 ], with ∆t = tk+1 − tk and normalized segment time s = (t − tk )/∆t ∈ [0, 1]. Action Matching. Action Matching [14] learns a scalar action network Sθ (t, x). For endpoint samples (x0 , x1 ) and an interpolated point xt = (1 − s)x0 + sx1 , the implemented segment loss has the form LAM = E w(tk )Sθ (tk , x0 ) − w(tk+1 )Sθ (tk+1 , x1 ) (13) + ∆t{w(t)(∂t Sθ (t, xt ) + 21 ∥∇x Sθ (t, xt )∥2 ) + Sθ (t, xt )∂t w(t)} . At test time, the velocity is vθ (t, x) = ∇x Sθ (t, x) and predictions are generated by solving ẋt = vθ (t, xt ). OT-CFM.
OT-CFM [15] pairs mini-batches by an optimal-transport plan πk∗ ∈ arg
min π∈Π(ρ̂k ,ρ̂k+1 )
E(x0 ,x1 )∼π ∥x0 − x1 ∥2 .
(14)
For paired endpoints, it uses the straight conditional path x1 − x 0 , (15) ∆t and trains a velocity network with the flow-matching objective E∥vθ (t, xt ) − ut ∥2 . The implementation uses the exact minibatch OT matcher from torchcfm and sets the conditional path noise to zero. xt = (1 − s)x0 + sx1 ,
ut =
MFM. Metric Flow Matching [34] uses a two-stage pipeline. It first trains a geodesic correction network ψη under a data-induced metric, then trains a velocity field using the corrected interpolant µt = (1 − s)x0 + sx1 + γ(t)ψη (x0 , x1 , t),
(16)
where γ(t) vanishes at the segment endpoints. The flow-matching target is the time derivative x1 − x 0 ut = + γ̇(t)ψη (x0 , x1 , t) + γ(t)∂t ψη (x0 , x1 , t), (17) ∆t with the last term present when the correction network depends explicitly on time. Our runs use the LAND-style data-manifold metric configuration. CURLY. CURLY [30] uses a related corrected interpolant, but with the normalized-time modulation s(1 − s): µt = (1 − s)x0 + sx1 + s(1 − s)ψη (x0 , x1 , s). (18) The target velocity used for the second-stage flow model is 1 [(x1 − x0 ) + (1 − 2s)ψη + s(1 − s)∂s ψη ] . (19) ut = ∆t The first stage fits ψη using pseudo-velocity supervision estimated from the observed temporal point clouds. Aligned CFM. Aligned CFM follows the adversarially learned interpolant approach of ALICFM [37]. It trains a global interpolant Iη (x0 , x1 , t) from the first to the last training snapshot. Intermediate marginals of Iη are aligned with observed training snapshots by an adversarial objective of the schematic form X min max [Ex∼ρ̂ℓ log Dℓ (x) + Ex0 ,x1 log{1 − Dℓ (Iη (x0 , x1 , tℓ ))}] . (20) η
D
ℓ
After this stage, the interpolant is frozen and a velocity model is trained with LALI−CFM = E∥vθ (t, Iη (x0 , x1 , t)) − ∂t Iη (x0 , x1 , t)∥2 . 23
(21)
DMSB. DMSB [29] is a stochastic bridge baseline in joint position-velocity state space zt = (xt , vt ). The implementation uses a momentum SDE discretization with learned control aθ (zt , t): vn 0 0 zn+1 = zn + ∆t + σ ∆t + , ϵn ∼ N (0, ∆tI). (22) 0 aθ (zn , tn ) ϵn Forward and backward policies are trained on the temporal grid, and the learned dynamics are rolled out to predict the held-out marginal distributions. DMSB appears only in the tables where this baseline was run.
C
PACE Implementation Details
PACE is implemented as a two-stage procedure. Stage 1 learns endpoint-preserving bridges and refines the cross-time matching; Stage 2 freezes the learned bridge correction and distills it into a global ODE velocity field for rollout. All reported experiments were run on a single NVIDIA A100 80GB GPU. Stage 1 anchor construction. For each experiment, the selected training frames are first grouped by time label. PACE stacks the training frames into an anchor tensor of shape K × N × d, where K is the number of training labels and N is the minimum selected cell count across training time points. If time points have unequal counts, each frame is trimmed to this common N for the Stage 1 matching problem. This produces one adjacent matching problem per training interval. Initial and refined matchings. The initial matching between adjacent anchors uses a cost combining normalized squared Euclidean distance with the local normal-direction penalty. After a fixed number of bridge-training epochs, PACE recomputes the matching by evaluating the learned path-action cost and solving the induced OT problem. We use the Python Optimal Transport (POT) package for OT computations. Rematching is repeated at a fixed experiment-specific interval (we usually choose 10 or 20), so bridge learning and correspondence estimation alternate throughout training. Bridge and flow distillation. For endpoints (x0 , x1 ) and local time s, the bridge has the endpointpreserving form γθ (x0 , x1 , s) = (1 − s)x0 + sx1 + s(1 − s)ψθ (x0 , x1 , s).
(23)
The bridge velocity ∂s γθ is obtained by automatic differentiation in 2D and by a Jacobian-vector product implementation in higher dimensions. After Stage 1, ψθ is frozen. Stage 2 samples paired endpoints from the Stage 1 matchings and trains a global velocity network vϕ (t, x) by regressing to the bridge velocity along generated points. Held-out predictions in the tables are produced by rolling out this Stage 2 velocity field.
D
Proof of Proposition 1 (Ill-posedness)
Proof of Proposition 1. For one interval [tk , tk+1 ], many couplings πk ∈ Π(ρ̂k , ρ̂k+1 ) can share the same endpoint marginals ρ̂k and ρ̂k+1 . For any such coupling, each coupled endpoint pair (x, y) can be connected by infinitely many C 1 paths γ : [0, 1] → Rd with γ(0) = x and γ(1) = y. Therefore, the observed snapshots determine neither a unique coupling nor a unique intermediate trajectory.
E
Local Metric and Correspondence
The metric used by PACE changes the matching criterion from endpoint proximity to path plausibility. In the locally constant idealization, suppose G = I + αPN ,
α > 0,
where PN is the orthogonal projector onto the local normal subspace. For a straight candidate displacement d = y − x, the metric action is d⊤ Gd = ∥d∥2 + α∥PN d∥2 . 24
(24)
Thus, two candidate matches with the same Euclidean distance are not equivalent: the one whose displacement is more tangent-aligned has lower action. The implemented initial matching uses this principle through a normal-projection penalty, while the refined matching applies the same idea to the learned curved bridge by averaging γ̇θ (s)⊤ G(γθ (s), t(s))γ̇θ (s)
(25)
over a finite probe grid. This is the practical reason PACE can prefer a slightly longer but tangentcompatible path over a shorter chord that cuts across the inferred developmental geometry.
F
Monotonicity of Ideal Alternating Updates
The alternating bridge-coupling updates can be viewed as block-coordinate descent for the discretized metric-action objective. For each adjacent time interval k, let Πk = Π(ρ̂k , ρ̂k+1 ) denote the feasible set of empirical couplings, and define the finite-dimensional objective J (θ, {πk }K−1 k=0 ) =
K−1 XX
(k)
πij cpath ij (θ),
(26)
k=0 i,j
where cpath ij (θ) is the discretized path-action cost in Eq. (9). This objective is the full-batch version of the metric-action term optimized during the bridge update, using the same path-action quadrature as the coupling update. Proposition 3 (Monotonicity under exact block updates). Assume that, at iteration r, the bridge update computes a global minimizer θr+1 ∈ arg min J (θ, {πkr }K−1 k=0 ), θ
and the coupling update computes, for each k, a global minimizer X (k) path πkr+1 ∈ arg min πij cij (θr+1 ). πk ∈Πk
i,j
Then the joint metric-action objective is non-increasing: r r K−1 J (θr+1 , {πkr+1 }K−1 k=0 ) ≤ J (θ , {πk }k=0 ).
(27)
Moreover, the decrease is strict whenever either the bridge block or at least one coupling block achieves a strict improvement over the previous iterate. Proof. By the optimality of the bridge update with {πkr } fixed, r r K−1 J (θr+1 , {πkr }K−1 k=0 ) ≤ J (θ , {πk }k=0 ).
By the optimality of each coupling update with θr+1 fixed, X (k),r+1 path X (k),r path πij cij (θr+1 ) ≤ πij cij (θr+1 ) i,j
for every k.
i,j
Summing these inequalities over k gives r+1 J (θr+1 , {πkr+1 }K−1 , {πkr }K−1 k=0 ) ≤ J (θ k=0 ).
Combining the two inequalities proves monotonicity. If either inequality is strict, the combined decrease is strict. (k)
Because Gk (x, t) = I + αCN (x, t) is positive definite, every path-action cost is nonnegative; therefore the sequence of ideal objective values is bounded below by zero and hence converges as a sequence of numbers. This does not imply convergence to a global optimum, nor does it imply uniqueness of the bridge or coupling. Equality can occur when a block update returns an equivalent minimizer or when the current block is already optimal. 25
Relation to the implemented algorithm. The proposition applies to the ideal metric-action problem in Eq. (26). The implemented PACE training approximates this scheme with stochastic minibatches, a neural bridge optimized by finite gradient steps, periodic rather than continuous rematching, and optional stabilizing regularizers. These choices are used for scalability and numerical stability, but they do not provide a per-gradient-step strict decrease guarantee for the realized neural training trajectory. The monotonic result should therefore be read as the optimization principle behind the alternating updates, while the empirical sections evaluate the behavior of the practical stochastic implementation.
G
Stage 1 Regularizers
PACE Stage 1 optimizes a metric action together with an optional regularization term: Lbridge = λmetric Lmetric + λreg Lreg .
(28)
The metric-action term is the discretized version of Eq. (8). The regularization term can include the implementation-level stabilizers below; in our experiments, this corresponds to a weighted combination λreg Lreg = λcoh Lcoh + λorth Lorth . Cross-segment velocity coherence. Let (xa , ta , va ) and (xb , tb , vb ) denote generated bridge points and velocities from different adjacent segments. PACE constructs Gaussian space-time weights ∥xa − xb ∥2 |ta − tb |2 Wab = exp − exp − 1{seg(a) ̸= seg(b)}, (29) σx,a σx,b σt,a σt,b where the bandwidths are estimated from local projector variation. The coherence loss penalizes nearby generated points from different segments when their normalized velocities disagree: P ⊤ va a,b Wab 1 − v̂a v̂b P Lcoh = , v̂a = . (30) W + ε ∥v ab a∥ + ε a,b Normal-motion suppression. For generated bridge points, PACE estimates a local normal projector and penalizes the normal component of the generated velocity: X 1 Lorth = ∥PN (x, t)v∥2 . (31) |G| (x,t,v)∈G
In 2D this reduces to a squared dot product with the locally estimated normal vector. In higher dimensions it uses the full normal-space projector PN = I − PT .
H
High-Dimensional Concentration Effects
The high-dimensional experiments in Section 4 operate in a regime where Euclidean probability mass and pairwise costs are strongly concentrated. This section explains the diagnostics used in Figure 4 and why concentration affects OT, MMD, and nearest-neighbor-based trajectory objectives. Norm concentration. For X ∼ N (0, Id ), the map f (x) = ∥x∥2 is 1-Lipschitz. The Gaussian concentration inequality therefore gives
and E∥X∥2 ≍
√
Pr(|∥X∥2 − E∥X∥2 | ≥ t) ≤ 2 exp(−t2 /2),
(32)
Std(∥X∥2 ) = O(d−1/2 ). E∥X∥2
(33)
d implies
This is the standard Lipschitz concentration phenomenon [39]: after whitening or PCA scaling, most points are pushed toward a thin shell rather than spreading across many radial scales. 26
Direction and nearest-neighbor concentration. For independent isotropic sub-Gaussian vectors X, Y ∈ Rd , normalized directions satisfy Y X , = Op (d−1/2 ), (34) ∥X∥2 ∥Y ∥2 so random directions become nearly orthogonal as d grows [40]. A related nearest-neighbor result states that if the relative variance of distances vanishes, Var(∥Q − X1 ∥2 ) → 0, (35) E[∥Q − X1 ∥2 ]2 then ! (d) (d) Dmax − Dmin Pr ≤ϵ →1 for every ϵ > 0, (36) (d) Dmin (d)
(d)
where Dmin and Dmax are the nearest and farthest distances from the query point [41]. Thus, when distance contrast collapses, the nearest neighbor and farthest neighbor become less distinguishable in relative terms. Effect on OT, MMD, and flow-matching losses. If two approximately isotropic samples X, Y ∈ Rd have coordinate-wise fluctuations with comparable scale, then ∥X − Y ∥22 =
d X
(Xj − Yj )2
(37)
j=1
is a sum of many coordinate-level contributions. Under standard independence or weak-dependence assumptions, a law-of-large-numbers or sub-Gaussian concentration argument implies that ∥X − Y ∥22 /d concentrates around its mean, with relative fluctuations that typically scale as O(d−1/2 ). For transport objectives, W22 (µ, ν) = inf E(x,y)∼π ∥x − y∥22 , (38) π∈Π(µ,ν)
so concentration directly reduces the dynamic range of the cost matrix used by OT and flow-matching couplings. For kernel metrics, MMD2 (µ, ν) = E k(x, x′ ) + k(y, y ′ ) − 2k(x, y) , (39) and a distance-based kernel such as k(x, y) = exp(−∥x − y∥22 /(2σ 2 )) also loses contrast when most ∥x − y∥22 values occupy a narrow interval. Consequently, different predicted distributions can receive similar numerical losses even when their local geometry differs. This is the metric-degeneration issue encountered by high-dimensional distance-based trajectory methods, including metric and OT flow-matching objectives [34]. Reference thresholds in Figure 4.
The first diagnostic is Std(∥X∥2 ) CVnorm = . (40) E∥X∥2 The horizontal reference level 0.3 is a practical warning threshold: below it, one standard deviation of radial variation is less than 30% of the mean radius, so most cells lie in a relatively thin shell and pairwise Euclidean costs are dominated by small angular or local fluctuations. This threshold is not a theorem-specific cutoff; it is a scale marker chosen to make the O(d−1/2 ) collapse visible on the empirical PCA representations. The second diagnostic compares time separation with within-time dispersion. For snapshot centers ci , cj , define ∥ci − cj ∥2 1 Rtime = mediani<j , r̄ij = EX∼µi ∥X − ci ∥2 + EY ∼µj ∥Y − cj ∥2 . (41) r̄ij 2 The reference level 1.0 marks the point where the median displacement between time-point centers is comparable to the typical within-time radius. When Rtime ≤ 1, the Euclidean shift between snapshots is no larger than the spread of the snapshots themselves, so time labels are difficult to separate using raw pairwise distances alone. Figure 4 summarizes these diagnostics in the main text. The appendix diagnostics in Figures A.6 and A.7 show the corresponding norm-concentration behavior for iPSC-Liu, OP-Cite, and OP-Multi. 27
Table A.8: iPSC-Liu [21] (10D/50D) per-timepoint results across holdouts t ∈ {4, 8, 12, 16}. MMD ↓
Dim Method t=4
W1 ↓
W2 ↓
t=4
t=8
t = 12 t = 16
t=4
t=8
t = 12 t = 16
10D Action Matching 0.6504 0.6959 0.5132 0.3626 3.8458 Aligned CFM 0.6245 0.5936 0.4051 0.2345 3.7937 CURLY 0.5461 0.6054 0.4936 0.5024 3.8694 DMSB 0.9250 0.7057 0.7448 0.7098 10.5837 MFM 0.5864 0.6674 0.5033 0.3521 3.5972 OT-CFM 0.6039 0.6339 0.4951 0.3251 3.6543
4.1608 3.4433 3.9773 4.1953 4.0239 4.0040
3.1014 4.0123 4.2076 2.5372 1.9877 4.2084 3.1387 3.4660 4.3490 5.5040 4.9522 10.6553 3.2515 2.7539 4.0205 3.2932 2.5884 4.0624
5.9396 4.0761 5.5607 4.5341 5.0654 5.5807
3.4279 2.9973 3.4974 5.6287 3.5968 3.6325
3.2507 2.3185 2.5294 3.7485
3.5755
2.7868 2.6406
PACE(ours)
t = 8 t = 12 t = 16
0.5244 0.5153 0.3113 0.2871 3.5918
3.9222 2.2482 3.4336 5.0190 2.6699 2.5030
50D Action Matching 0.3517 0.3506 0.3278 0.2197 7.2261 8.4554 7.9593 6.8096 7.8273 10.3927 8.5704 7.1567 Aligned CFM 0.3519 0.2917 0.2397 0.1663 7.0562 7.5723 7.5947 7.0302 7.6855 8.2791 8.2483 8.0080 CURLY 0.3824 0.4189 0.2235 0.2327 8.2766 9.7010 7.6085 6.8715 8.8476 11.2946 8.2990 7.1972 DMSB 0.7593 0.5730 0.4787 0.4441 14.5899 10.4588 9.7531 8.9456 14.7121 10.7933 10.1192 9.3447 MFM 0.3539 0.3467 0.2679 0.1728 7.0197 9.2505 8.5746 6.3246 7.6665 10.5394 9.2502 6.7254 OT-CFM 0.3530 0.3458 0.2802 0.1803 7.1278 9.0608 8.7395 6.4288 7.7653 10.4423 9.3687 6.8054 PACE(ours)
0.3609 0.3004 0.2193 0.1396 7.0656
6.8194 6.8718 6.2581 7.6896
7.0798
7.5736 6.7914
Table A.9: OP-Cite and OP-Multi [36] (100D) results. Method
OP-C ITE (100D)
OP-M ULTI (100D)
MMD ↓
W1 ↓
W2 ↓
MMD ↓
W1 ↓
W2 ↓
Action Matching Aligned CFM CURLY MFM OT-CFM
0.1753 0.1434 0.2529 0.1435 0.1448
9.9388 10.7009 11.5163 10.4301 10.4447
10.0448 10.8094 11.6015 10.5409 10.5491
0.1622 0.1720 0.2786 0.1699 0.1662
10.7711 10.8330 12.0015 10.6274 10.6128
10.8160 10.8826 12.0440 10.6702 10.6564
PACE(ours)
0.1409
10.4574 10.5605
0.1693
10.5824 10.6293
High-dimensional benchmark results. Table A.8 reports the per-timepoint reconstruction metrics for iPSC-Liu [21] in 10D and 50D PCA representations, evaluated on holdout time points t ∈ {4, 8, 12, 16}. Table A.9 gives the corresponding results for the multimodal OP-Cite and OP-Multi datasets [36] in 100D. These tables complement the time-averaged summary in Table 3 and the concentration diagnostics in Figure A.6.
data KDE mean = 67.55
0.015 0.010 0.005 0.000
0.10
data KDE mean = 72.83
0.06
0.020 density
density
0.020
µ = 72.830 σ = 18.979 CV = 0.2606
0.025 0.015 0.010
100
kxk 2
150
op-cite (d = 100)
0.000
data KDE mean = 20.00
0.04 0.02
0.005 50
µ = 20.001 σ = 7.223 CV = 0.3611
0.08 density
µ = 67.545 σ = 18.838 CV = 0.2789
0.025
50
100
kxk 2
150
op-multi (d = 100)
0.00
20
40
kxk 2
60
80
iPSC (d = 50)
Figure A.6: Norm-concentration diagnostics for iPSC-Liu [21] and OP-Cite/OP-Multi [36] representations. The panels illustrate the empirical thin-shell behavior summarized by CVnorm , with the 0.3 reference level used as a practical concentration warning threshold.
I
iPSC-Liu Results
Figure A.7 shows the norm-concentration diagnostics for iPSC-Liu across increasing PCA dimensions, illustrating the O(d−1/2 ) norm-concentration behavior discussed in Appendix H. These panels correspond to the high-dimensional benchmark experiments reported in Table A.8 and summarized in the main text in Table 3. 28
data KDE mean = 13.10
10
20
kxk 2
iPSC (d = 2)
30
0.12 0.10 0.08 0.06 0.04 0.02 0.00
µ = 16.704 σ = 6.575 CV = 0.3936
0.10
data KDE mean = 16.70
0.06
µ = 20.001 σ = 7.223 CV = 0.3611
0.08 density
µ = 13.097 σ = 5.462 CV = 0.4171
density
density
0.12 0.10 0.08 0.06 0.04 0.02 0.00 0
data KDE mean = 20.00
0.04 0.02
20
kxk 2
40
iPSC (d = 10)
60
0.00
20
40
kxk 2
60
80
iPSC (d = 50)
Figure A.7: Norm-concentration diagnostics for iPSC-Liu [21] representations. Increasing PCA dimension reduces relative radial variation, consistent with the O(d−1/2 ) norm-concentration behavior described in Appendix H.
J
Intuition behind alternating bridge and coupling optimization
Why coupling matters for the bridge. The neural bridge is trained on endpoint pairs sampled from the current coupling. If the coupling is biologically implausible (for example, pairing an early stem cell with a terminally differentiated cell), the bridge must learn a path that traverses the entire developmental manifold in a single segment. Such a path is likely to pass through regions where the metric strongly penalizes normal motion, yielding high action and making it difficult for the network to find a low-action deformation. Conversely, when the coupling pairs cells at nearby pseudotime stages, the displacement is small and largely tangent-aligned, so the bridge can easily learn a geometry-consistent interpolant. In short, the coupling determines the training set for the bridge; wrong pairs give the network an impossible learning problem. Why the bridge matters for the coupling. The coupling update solves an OT problem whose cost matrix is the path action evaluated on the neural bridge (Eq. (9)). If the bridge is untrained, its paths are nearly straight lines and the cost reduces to Euclidean distance; the OT solver then cannot distinguish a biologically plausible pairing from an implausible one. Once the bridge has learned metric-aware paths, the action cost becomes geometry-sensitive. Pairs whose displacement aligns with the local developmental tangent incur low action, while pairs that would require motion across normal (non-developmental) directions incur high action. The bridge therefore transforms an uninformative Euclidean cost into a geometry-aware cost, allowing OT to select biologically meaningful pairings. Why amortize with a shared neural network. Without amortization, every candidate endpoint pair would require solving an independent boundary-value problem (the geodesic equation) under the state-dependent metric Gk . With Nk Nk+1 candidate pairs, this is computationally infeasible. The neural bridge amortizes this cost by learning a single parametric family γθ (x, y, τ ) that approximates the minimum-action path for all pairs. The approximation need not be perfect at initialization; it only needs to be good enough to provide a more informative cost than Euclidean distance, and it improves as the coupling improves.
K
Adaptive tangent and normal projectors
The weighted local covariance at each anchor point yields eigenvalues λ1 ≤ λ2 ≤ · · · ≤ λd and orthonormal eigenvectors v1 , . . . , vd . The tangent subspace is spanned by the leading (maximumvariance) eigenvectors. Its effective dimension qr is chosen as the smallest integer such that the cumulative explained variance reaches a prescribed threshold τ ∈ (0, 1) (e.g., 0.95): ( ) Pq j=1 λd−j+1 qr = min q : ≥τ . Pd j=1 λj 29
(r)
Pqr
⊤ j=1 vd−j+1 vd−j+1 and the normal projector is its orthogonal (r) (r) (r) complement PN = I − PT . In two dimensions this reduces to PN = nr n⊤ r with nr = v1 and qr = 1.
The tangent projector is PT
L
=
Adaptive bandwidth estimation (k)
(k)
For each segment k, the spatial and temporal bandwidths hx and ht are estimated from the anchor normal bank rather than set by hand. The estimation aggregates information over a small window of anchor snapshots centered on segment k (by default one snapshot on each side). The same local (r) eigendecomposition also yields the tangent subspace PT (Appendix K), which is used for bandwidth estimation below but does not enter the metric tensor directly. Spatial geometry per snapshot. For each anchor snapshot a in the window, PACE computes three quantities from the k nearest neighbors of every cell (excluding the cell itself): (a)
• Local spatial spacing ∆x : the median tangential step size ∥⟨xj − xi , ti ⟩∥ between each cell i and its neighbors j, where ti is the local tangent direction. This captures how densely cells are spaced along the developmental manifold. (a)
(j)
(i)
(a)
• Normal spatial rate ρN,x : the median Frobenius-norm difference ∥PN − PN ∥F /∆x between the normal projectors of neighboring cells. This measures how rapidly the local geometry changes in space. (a)
• Tangent spatial rate ρT,x : the analogous quantity for tangent projectors. Cross-snapshot temporal rates. For each consecutive pair (a, a+1) in the window, PACE matches cells by nearest-neighbor correspondence and computes: (a,a+1)
(nn(i))
(i)
• Normal temporal rate ρN,t : the Frobenius-norm difference ∥PN −PN ∥F /|ta+1 − ta | between matched cells, divided by the experimental time gap. This measures how rapidly the local geometry evolves over time. (a,a+1)
• Tangent temporal rate ρT,t
: the analogous quantity for tangent projectors.
Segment bandwidth assembly. All per-snapshot and cross-snapshot quantities are concatenated (k) (k) (k) (k) and their positive medians are taken, yielding four segment-level scalars: ∆x , ρN,x , ρN,t , and ρT,t . The metric bandwidths are then ! 1 (k) h(k) , (42) x = max ∆x , (k) ρN,x (k) ρN,t (k) cN = (k) , ρN,x
(43)
(k)
(k)
ht
=
hx
(k)
.
(44)
cN
The spatial bandwidth hx is the larger of the local spacing and the reciprocal spatial variation scale, preventing the kernel from being narrower than either the point density or the geometry variation. (k) The ratio cN converts spatial scale into equivalent temporal scale through the speed at which the (k) (k) normal geometry changes in time relative to space. A similar pair (σx , σt ) is computed from the (k) (k) tangent rates for the coupling-refinement kernel, but the metric field uses (hx , ht ). 30
M
Limitations and Future Work
PACE has several limitations. First, its local metric depends on nearest-neighbor covariance estimates within each snapshot. When data are sparse, noisy, or highly heterogeneous, the estimated tangent and normal subspaces may be unstable; adaptive neighborhood selection may improve robustness. Second, PACE is optimized with stochastic neural training and periodic rematching. Although Appendix F states a monotonicity property for ideal exact block-coordinate updates, the implemented algorithm is only a scalable approximation and does not guarantee global optimality. Third, destructive snapshot data do not provide ground-truth correspondences or continuous cell histories. The reported metrics evaluate held-out marginal reconstruction rather than individual trajectory correctness. PACE-inferred paths should be interpreted as plausible geometry-consistent reconstructions, not directly observed cellular histories. Future work will extend PACE beyond closed-population transport by incorporating birth-death or growth terms [3]. Another direction is to learn representations that preserve local trajectory geometry in high-dimensional settings, where distance concentration can make coupling costs less discriminative.
31