Technical Summary: PACE – Geometry-Aware Bridge Transport for Single-Cell Trajectory Inference

1. Problem Formulation

Single-cell trajectory inference aims to reconstruct continuous cell-state evolution from destructive time-course snapshots. A fundamental challenge in this domain is that the problem is ill-posed: neither cross-time cell correspondences nor the continuous paths between snapshots are observed. Consequently, observed snapshot distributions alone do not uniquely determine the underlying dynamics.

Existing optimal transport (OT) and flow-based methods typically estimate couplings based on Euclidean proximity at observed clock times. This approach assumes that cells sampled at the same experimental time occupy similar developmental stages. However, in asynchronous biological processes (e.g., reprogramming or differentiation), cells at the same clock time may reside at different latent pseudotime stages or distinct fate branches. Relying solely on Euclidean distance can lead to misaligned trajectories, coupling cells from incompatible developmental programs (e.g., linking apoptotic cells to rapidly expanding stem cells) or failing to capture curved, anisotropic motion along the cell-state manifold.

2. Methodology: PACE Framework

The authors propose PACE (Geometry-Aware Bridge Transport), a framework that selects geometry-consistent continuous transport dynamics from unpaired snapshots. The core intuition is that while true cell identities are unobserved, local spatiotemporal neighborhoods provide a weak prior on admissible motion directions: motion along locally supported tangent directions is more plausible than motion orthogonal to the observed manifold structure.

PACE operates through three coupled components:

A. Time- and State-Dependent Anisotropic Riemannian Metric

PACE replaces fixed Euclidean costs with a Riemannian path action induced by a metric tensor $G_k(x, t)$.

• Local Geometry Estimation: For each observed cell, PACE estimates local tangent and normal subspaces using nearest-neighbor covariance analysis within the same snapshot.
• Metric Construction: The metric is defined as $G_k(x, t) = I + \alpha C^{(k)}_N(x, t)$, where $C^{(k)}_N$ is a spatiotemporally interpolated field of local normal projectors.
• Cost Function: The cost of a path $\gamma$ is the integral of the velocity squared weighted by the metric:
  $$A_{G_k}[\gamma] = \int_0^1 \dot{\gamma}(\tau)^\top G_k(\gamma(\tau), t_{k,\tau}) \dot{\gamma}(\tau) d\tau$$
  This formulation penalizes velocity components in the locally estimated normal direction (high cost) while preserving low cost for motion along tangent directions.

B. Alternating Bridge and Coupling Optimization

To solve the resulting infinite-dimensional transport problem, PACE employs a finite-dimensional alternating optimization procedure:

1. Bridge Learning (Fixed Coupling): Given a cross-time coupling $\pi_k$, PACE trains an endpoint-preserving neural bridge $\gamma_\theta(x, y, \tau)$. The bridge is parameterized as a straight line plus a neural correction term scaled by $\tau(1-\tau)$, ensuring $\gamma_\theta(x, y, 0)=x$ and $\gamma_\theta(x, y, 1)=y$. The network is trained to minimize the anisotropic path-action loss defined by the metric $G_k$.
2. Coupling Refinement (Fixed Bridge): Using the learned bridge, PACE computes the path-action cost between all candidate endpoint pairs. It then solves an Optimal Transport problem to update the coupling $\pi_k$, minimizing the total path action rather than Euclidean distance.
   This creates a bootstrap mechanism where better couplings provide better endpoint supervision for the bridge, and better bridges provide geometry-aware costs for refining the coupling.

C. Velocity Field Distillation

The endpoint-conditioned bridges are pair-specific. To enable global trajectory inference, PACE distills the learned bridge dynamics into a global continuous-time velocity field $v_\phi(x, t)$. This field is trained to match the velocities of the neural bridges at arbitrary states and times, allowing for the generation of continuous trajectories via ODE integration from any initial distribution.

3. Key Contributions

1. Geometry-Aware Metric: Introduction of a time- and state-dependent Riemannian metric based on local spatiotemporal tangent-space projections. This provides a path-action cost that selects couplings beyond simple Euclidean endpoint proximity.
2. Iterative Bridge-Transport Procedure: Development of an alternating optimization scheme that jointly refines cross-time correspondences and interpolant paths under the same geometry-aware action, avoiding the need for fixed endpoint pairings.
3. Global Velocity Distillation: A method to distill endpoint-conditioned neural bridges into a global velocity field for continuous trajectory generation without explicit cell pairing or lineage tracing.

4. Experimental Results

PACE was evaluated on seven datasets (controlled synthetic and biological) covering nine held-out reconstruction experiments.

• Controlled 2D Benchmarks: On the "Ocean" (rotational non-gradient) and "Two-Branch" benchmarks, PACE achieved the strongest overall reconstruction performance. It successfully recovered smooth rotational motion and branching structures without velocity supervision.
• Biological Benchmarks:
  • EB PHATE: PACE improved held-out marginal reconstruction (MMD, W1, W2) and showed a 15.4% improvement in RNA-velocity alignment compared to baselines, despite RNA velocity not being used during training.
  • iPSC-Liu & Schiebinger2019: PACE consistently outperformed baselines (Action Matching, Aligned CFM, CURLY, DMSB, MFM, OT-CFM) across differentiation and reprogramming tasks.
• Quantitative Gains: Across all datasets, PACE reduced MMD, W1, and W2 by an average of 23.7% relative to the strongest competing baseline.
• High-Dimensional Performance: In high-dimensional settings (10D, 50D, 100D PCA), where Euclidean distances suffer from concentration effects, PACE maintained robust performance by leveraging local anisotropy rather than global separation.
• Ablation Studies: Removing components such as coupling rematching, the anisotropic metric penalty ($\alpha=0$), or spatial/temporal kernel smoothing degraded performance, confirming that the gains arise from the combination of these elements.

5. Significance and Claims

The paper claims that local geometry serves as a useful inductive bias for reconstructing trajectories from destructive snapshots. By replacing Euclidean costs with a geometry-aware path action, PACE addresses the fundamental ambiguity of trajectory inference without requiring:

• Explicit cell pairing or lineage tracing.
• RNA velocity supervision during training.
• Assumptions of synchronous development.

The authors emphasize that PACE provides plausible, geometry-consistent reconstructions rather than ground-truth cellular histories, as the underlying problem remains ill-posed. The method demonstrates that inferring admissible directions from local spatiotemporal neighborhoods is sufficient to significantly improve the accuracy of population-level dynamics inference in single-cell biology.