Technical Summary: LoCoFlow

Problem Statement

Single-cell sequencing provides high-resolution snapshots of cellular states at discrete time points, but the destructive nature of the measurement prevents the observation of true longitudinal correspondences between individual cells. Reconstructing continuous developmental trajectories from these snapshots is an ill-posed problem with multiple possible solutions. Furthermore, existing methods often treat cells as isolated vectors, failing to adequately represent the local manifold geometry of cell-state space or the inherent uncertainty in lineage diversification. While cells with similar transcriptional states may exhibit divergent fate tendencies, current optimal transport and flow-matching approaches typically infer temporal couplings at the individual cell level, lacking a robust framework to model the evolution of local cell neighborhoods as cohesive units.

Methodology: LoCoFlow

The authors introduce LoCoFlow (Local-Context Flow Matching), a framework that integrates local point-cloud geometry with conditional fate uncertainty to model single-cell dynamics. The method operates in two primary stages:

1. Local Patch Construction and Hierarchical Correspondence

Instead of matching isolated cells, LoCoFlow constructs overlapping "patches" (local neighborhoods) around anchor cells.

• Patch Definition: For expression data, patches are geodesic neighborhoods on a nearest-neighbor graph; for spatial data, they are fixed-size neighborhoods in measured coordinates.
• Hierarchical Transport: The method employs a two-level optimal transport (OT) scheme. First, it solves a patch-level transport problem using Bures distance to align patch centers and covariances (capturing both location and dispersion). Second, within retained patch pairs, it computes an entropy-regularized cell-level transport plan.
• Soft Assignment: The aggregation of these plans yields a conditional distribution of target cells for each source, rather than a hard assignment, preserving multiple potential successor correspondences.

2. Two-Stage Training Pipeline

• Stage 1: Geometry-Regularized Bridge Transport:
  • Using the paired endpoint sets derived from the hierarchical transport, the model learns "bridges" (intermediate paths) between source and target patches.
  • A neural network predicts a residual velocity field $f_\theta$ that, combined with linear interpolation, ensures the path passes exactly through the endpoints.
  • Geometric Priors: The loss function includes kinetic energy terms and a radial penalty that constrains relative deformation (expansion/contraction) of the patch without restricting common translation. It also penalizes velocities orthogonal to the estimated local manifold geometry.
• Stage 2: Conditional Velocity Distillation:
  • Since the future state is unknown during inference, the bridge (which requires both endpoints) cannot be used directly. LoCoFlow distills the bridge velocities into a conditional velocity field $V_\omega$.
  • Hierarchical Latents: The model learns to generate target codes $c_i$ from a source context $Z_0$ using a hierarchical latent structure: a patch-shared latent $\eta_P$ (capturing population-level outcomes) and cell-specific latents $\eta_i$ (capturing individual variation).
  • Training: A target encoder first learns to map observed endpoints to codes $c^\star$. Then, a generator learns to produce these codes from the source context and latents. The velocity network is trained to match the bridge velocities conditioned on these generated codes.
  • Inference: At inference, latent codes are drawn from a prior (without future observations) and held fixed during integration. This allows for the sampling of alternative trajectories representing fate uncertainty.

Key Contributions

1. Local-Context Representation: Shifts the fundamental unit of dynamic modeling from isolated expression vectors to local population states (point clouds), explicitly modeling the evolution of neighborhood geometry.
2. Hierarchical Transport: Introduces a coupled patch-level and cell-level optimal transport mechanism that respects both the global motion of cell populations and their internal structural changes (e.g., heterogeneity, expansion).
3. Conditional Fate Uncertainty: Develops a two-stage distillation process that separates geometric path constraints from outcome uncertainty, enabling the sampling of multiple plausible fate trajectories without future observations.
4. Perturbation Analysis Framework: Provides a mechanism to distinguish between focal-cell effects and neighborhood (background) effects in gene perturbation simulations, revealing how local context shapes predicted responses.

Experimental Results

The paper evaluates LoCoFlow across three datasets: embryoid body (EB) differentiation, iPSC reprogramming, C. elegans development, and Axolotl brain regeneration.

• Held-Out Reconstruction: On EB and iPSC datasets, LoCoFlow achieves lower Wasserstein distances ($W_1, W_2$) compared to baselines like OT-CFM, MFM, and PACE in most high-dimensional settings (e.g., PCA50). However, the paper notes that LoCoFlow sometimes exhibits higher Maximum Mean Discrepancy (MMD), suggesting predictions may be more concentrated in specific regions of the cell-state space.
• Fate Prediction (LARRY Data): Using lineage-traced LARRY data, LoCoFlow reduces the mean raw Brier error from 0.7990 (scDiffEq) to 0.6620, indicating improved accuracy in predicting full fate-probability vectors.
• Perturbation Responses:
  • Geometric Patches: When perturbing genes in geometrically defined patches, background-only and all-member interventions produce aligned responses (e.g., shifting fate from monocytes to neutrophils).
  • Random Minibatches: In contrast, models using randomly assembled patches show opposing focal-cell and background responses, highlighting the importance of geometric context.
• Spatial Interpolation (Axolotl): LoCoFlow improves spatial prediction ($W_2$) for Axolotl brain regeneration compared to OT-CFM, though the authors note that predicted spatial distributions remain visually similar to the source time point, suggesting a need for additional spatial alignment.

Significance and Claims

The paper claims that LoCoFlow provides a robust framework for studying how local cell population structure shapes inferred trajectories while explicitly representing uncertainty in cell fate. The authors emphasize that:

• Local Context is Critical: The distinct response patterns observed between geometric patches and random minibatches support the hypothesis that local context is a useful representation for probabilistic fate inference.
• Modeling Assumptions: The authors maintain a modest stance, acknowledging that snapshot data do not uniquely identify cellular histories. The geometric and latent components of LoCoFlow impose modeling assumptions rather than revealing true lineage relationships.
• Testable Hypotheses: The framework generates testable hypotheses regarding how local context shapes predicted responses to perturbations, distinguishing between intrinsic cell changes and environmental influences.

The work does not claim to solve the fundamental ambiguity of trajectory inference but offers a method that better aligns with the manifold geometry of cell states and provides a structured way to quantify and sample fate uncertainty.