PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yuhao Sun, Zekun Wu, Zixun Huang, Peijie Zhou

arXiv:2608.21070v1cs.LGcs.AIq-bio.GN

TL;DR

Sparse temporal observations and first-order trajectory models limit inference of complex biological dynamics and lineage structure. TracingFlow addresses this with simulation-free second-order Flow Matching that regresses acceleration and initial velocity for DOAT. On simulated and real-world data, it reports improved distribution reconstruction and preservation of biological priors, while retaining computational and modeling limitations.

  • Problem

    Most trajectory-inference frameworks assume first-order dynamics, while practical single-cell data provide positions without velocities, limiting second-order modeling of biological dynamics.

  • Method

    TracingFlow solves DOAT through simulation-free second-order Flow Matching by regressing acceleration fields and initial velocity, with an iterative strategy for VM-DOAT.

  • Results

    TracingFlow achieves competitive or improved reconstruction accuracy and preserves biological priors on multi-time-point real-world datasets, with lineage-prior experiments showing improved consistency.

  • Takeaways & Limitations

    Second-order dynamics provide greater expressivity for temporal distribution recovery, while biological-prior integration better preserves intercellular lineages.

  • Takeaways & Limitations

    SOAT pre-calculation is computationally intensive, and the method approximates conditional velocity distributions using their expectation because explicit modeling is prohibitive.

Abstract

from arXiv · show

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

1 Introduction

Trajectory Inference recovers continuous dynamics from discrete observations, but prevailing methods often use first-order dynamics that limit expressiveness. TracingFlow addresses this with simulation-free second-order modeling, acceleration regression, and biological-prior integration.

  • Trajectory Inference recovers underlying dynamics from discrete observations in single-cell omics, with Optimal Transport identifying least-cost evolution between distributions.
  • Most trajectory-inference frameworks assume first-order dynamics, restricting velocity to a single value for each position and time.This constraint can limit modeling of complex biological priors such as lineage tracing.
  • Second-order alternatives address richer dynamics but may iteratively retrain acceleration fields and estimate initial velocity heuristically rather than optimizing it jointly.
  • TracingFlow formulates Dynamical Optimal Acceleration Transport and directly regresses acceleration fields and initial velocity without ODE simulation.The framework is designed as a simulation-free second-order trajectory-inference method.
  • TracingFlow demonstrates competitive or improved reconstruction accuracy and biological-prior preservation on multi-time-point real-world datasets.

2 Related Works

Flow Matching offers scalable simulation-free transport, but related approaches commonly use first-order dynamics and struggle with intersecting trajectories or complex lineage structure. TracingFlow extends continuous flow matching with second-order dynamics and embedded lineage priors.

  • Flow Matching provides scalable simulation-free training by learning flow maps that transport probability distributions.
  • Most Flow Matching variants use first-order dynamics, while second-order approaches may require iterative acceleration training and heuristic initial-velocity estimation.
  • Standard Flow Matching’s velocity dependence on only position and time prevents representing intersecting trajectories with distinct velocities.Prior work adds class labels, initial positions, or hidden states to decouple velocity from current position.
  • TracingFlow embeds lineage-tracing priors into a continuous second-order OT flow-matching framework rather than relying only on discrete couplings or static matrices.

3 Mathematical Background

DOAT models particle position and velocity jointly under second-order dynamics, constraining distributions at multiple timestamps. VM-DOAT captures the practical setting where only positions are observed and velocities are unknown.

  • DOAT defines an augmented state s = (x, v) with dynamics ˙x = v and ˙v = a(x, v, t).
  • The distributional formulation evolves a density through a continuity equation in position–velocity space while matching prescribed joint densities at K + 1 timestamps.
  • DOAT minimizes acceleration-based transport cost over these second-order dynamics and assumes an optimal flow map for particle-based numerical solution.
  • VM-DOAT addresses practical single-cell datasets where only particle-position marginals are observed and velocity quantities are unavailable.
  • VM-DOAT imposes constraints only on position marginals, leaving the velocity distribution unconstrained.

4 Methodology of TracingFlow

TracingFlow solves DOAT by combining optimal coupling, cubic single-particle trajectories, and acceleration flow matching, then bridges missing-velocity observations through a global completion scheme. Its constructed paths recover prescribed marginals and converge to DOAT solutions under vanishing noise.

  • Method overview: TracingFlow decomposes DOAT into optimal coupling and single-particle optimal trajectories, then uses global path-based velocity completion to address VM-DOAT.
  • Optimal single-particle trajectory: The unique minimum-acceleration trajectory between boundary positions and velocities is a cubic polynomial.
  • Optimal coupling: The resulting particle cost defines the optimal coupling problem between endpoint states in augmented position–velocity space.
  • Optimal coupling: DOAT permits crossings in projected position space but prohibits trajectory crossings in augmented space because intersections imply strictly higher transport cost.
  • Conditional and marginal probability path: Conditional probability paths mix trajectories selected from adjacent-time optimal plans, and the marginal path recovers the true distribution at each prescribed time point as σx, σv → 0.
  • Acceleration flow matching: Conditional Acceleration Flow Matching trains on simpler conditional fields while preserving the optimization gradients of marginal acceleration flow matching up to a parameter-independent constant.
  • Conditional and marginal probability path: The marginal probability path and acceleration field converge to the DOAT solution as σx and σv approach zero.
  • VM-DOAT to DOAT: Missing velocities are obtained by solving a sparse tridiagonal linear system in O(K) time, transforming VM-DOAT into a solvable DOAT instance.

5 TracingFlow Algorithm for Trajectory Inference

TracingFlow approximates conditional velocities, solves DOAT with flow matching, and incorporates biological lineage priors into transport costs. Its minibatch strategy improves computational efficiency, while theoretical guarantees relate reconstruction error to fitting losses.

  • 5.1 Approximation of Conditional Velocity Distribution: TracingFlow deterministically assigns one velocity to each data point because estimating the full conditional velocity distribution is computationally prohibitive.This approximation transforms VM-DOAT into a DOAT problem before learning the acceleration field.
  • 5.2 MiniBatch-OT: The minibatch strategy partitions consecutive empirical distributions into corresponding minibatches, significantly enhancing computational efficiency.The strategy exploits the Dirac-delta approximation of conditional velocity distributions.
  • 5.3 Incorporation of Biological Priors: Lineage tracing priors modify the transport cost matrix by multiplying cross-barcode costs by a penalty factor p0 > 1.This soft constraint discourages transitions between distinct lineages during dynamics learning.
  • 5 TracingFlow Algorithm for Trajectory Inference: TracingFlow computes optimal transport plans, samples conditional probability paths, and regresses the acceleration network aθ(x, v, t).Initial velocities are estimated iteratively before the transport plan and acceleration network are trained.
  • Theoretical Guarantees: When acceleration and initial velocity are perfectly fitted, TracingFlow recovers the exact distribution; otherwise, positional W2 error is bounded by velocity and AFM losses.The bound provides a theoretical link between model-fitting losses and distributional reconstruction error.

6 Experiment Results

Experiments evaluate transport, interpolation and extrapolation, and lineage-aware inference across synthetic and single-cell datasets. TracingFlow learns crossing trajectories, achieves the highest held-out interpolation accuracy, and outperforms other algorithms when biological priors are incorporated.

  • Distribution Transport: TracingFlow yields lower average W1 and W2 distances across evaluated datasets than first-order algorithms.The evaluation covers a 2D synthetic dataset and real single-cell omics datasets of varying dimensions.
  • Distribution Transport: TracingFlow learns crossing trajectories in position space, whereas OT-CFM learns non-crossing paths on the 2D Simulation dataset.The comparison illustrates the additional trajectory patterns enabled by the second-order model.
  • Interpolation and Extrapolation: TracingFlow achieved the highest accuracy in the EB 5D Hold-One-Out experiment for interpolated held-out distributions.The result supports second-order dynamics for capturing high-curvature and nonlinear trajectories.
  • Lineage Tracing Data: With lineage tracing priors, TracingFlow outperformed other algorithms on lineage-weighted W1 and W2 distances.On the Hematopoiesis dataset, its inferred positions for two barcodes were significantly closer to ground truth than OT-CFM’s.

7 Conclusion and Limitation

TracingFlow uses second-order dynamics to recover temporal distributions and preserve biological lineage structure. Its main limitations are computational SOAT pre-calculation and the lack of a clear physical interpretation for the acceleration-cost objective.

  • Conclusion: TracingFlow enables efficient, simulation-free solutions for DOAT by learning acceleration and initial velocity, while second-order dynamics enhance expressivity and temporal distribution recovery.The framework is described as preserving intercellular lineages when biological priors are incorporated.
  • Limitation: SOAT pre-calculation is computationally intensive, although minibatch-OT can mitigate this limitation.
  • Conclusion: On Sim-Lineage, the comparison distinguishes OT-CFM paths without biological priors from TracingFlow paths that preserve them.
  • Limitation: The acceleration-cost objective lacks a clear physical interpretation because standard classical actions exclude second time derivatives.

A.6 Proof of Proposition 4.7

The proof establishes that TracingFlow’s probability path converges to the optimal DOAT solution as smoothing vanishes, while learned-flow distribution error admits a Wasserstein bound tied to training losses.

  • DOAT optimality: As σx and σv approach zero, TracingFlow’s constructed path exactly solves the DOAT problem.
  • Optimal trajectories: The minimum single-particle acceleration-cost trajectory connecting endpoint positions and velocities is a cubic polynomial.
  • Approximation guarantee: The positional W2 distance between generated and true distributions is bounded using acceleration-field and initial-velocity losses.
  • Evaluation: Lineage-weighted W1 and W2 are indicators of lineage-prior preservation rather than mathematically symmetric distance metrics.

B.3 Datasets

The evaluation uses synthetic dynamics and several real single-cell datasets spanning dimensions, time points, cell counts, and lineage information.

  • Simulation datasets: The 3D simulation lineage dataset contains three lineage barcodes, a twist structure, three time points, and 500 particles sampled per initial condition.
  • Real datasets: The Cite-seq dataset contains 31,240 cells across four time points and is reduced to 100- and 5-dimensional PCA representations.
  • Real datasets: The Embryoid Bodies dataset contains 16,819 cells sampled at five time points and represented in five PCA dimensions.
  • Real datasets: The Hematopoiesis dataset contains 130,887 cells across three time points, with 22,329 barcode-labelled cells selected for 50-dimensional PCA analysis.
  • Real datasets: The Mexico Gulf dataset is a two-dimensional dataset with nine time points used to assess long-term system-dynamics capture.

C.1 Experiment Details across various datasets

Across synthetic, held-out-time, lineage, and real-world datasets, TracingFlow generally reconstructs distributions accurately and preserves biological priors. Its second-order formulation also produces smooth trajectories, with higher training time than first-order methods.

  • Distribution reconstruction: TracingFlow reconstructs the 2D Simulation distribution with higher average accuracy than other algorithms.Table 4 reports W1 and W2 distances at each time point.
  • Distribution reconstruction: TracingFlow achieves superior reconstruction accuracy on Cite 5D and Cite 100D at every time step except t = 3.The corresponding learned trajectories are shown in Figure 6.
  • Interpolation: TracingFlow outperforms other methods on EB 5D held-out-time reconstruction at every time point except t = 1.Figure 7 visualizes the interpolated distributions across four held-out time points.
  • Biological-prior preservation: TracingFlow obtains lower lineage-weighted W1 and W2 scores on the 3D Simulation Lineage dataset because its trajectories align with biological priors.This indicates stronger preservation of the supplied lineage structure.
  • Biological-prior preservation: On Hematopoiesis, TracingFlow surpasses other algorithms at all time points with biological priors and maintains superior Vanilla W1 and W2 performance without them.Figure 8 visualizes generated data under both conditions.
  • Additional evaluation: TracingFlow achieves the best distribution reconstruction accuracy at all time points on Mexico Gulf while yielding smooth trajectories.Second-order methods generally require more training time than first-order methods, although TracingFlow is slightly faster than 3 MSBM.

C.3 Sensitivity Analysis on Minibatch-OT

The minibatch-OT analysis evaluates whether TracingFlow can reduce the computational burden of full optimal transport while retaining reconstruction accuracy. Within a batch-size range around 10^3, reconstruction remains consistent and accurate.

  • Motivation: Full-dataset OT plans become rapidly more expensive as data scale increases, motivating minibatch computation.The experiments use minibatch-OT sensitivity analysis on Cite-5D to assess this trade-off.
  • Sensitivity result: TracingFlow consistently and accurately reconstructs distributions when minibatch size is varied within a range around 10^3.These results are reported in Table 13.
  • Training procedure: TracingFlow training assigns velocities through iterative SOAT, computes transport couplings and propagators, and then trains acceleration and initial-velocity networks.The procedure includes acceleration matching and initial velocity losses.
  • Method comparison: Compared with related methods, 3 MSBM uses iterative coupling and acceleration-field training, whereas MMFM retains a velocity field single-valued in x.OAT-FM fine-tunes a pretrained flow-matching model rather than solving the full multi-marginal augmented-space problem.

E.2 Is S = X × V a Phase Space?

The paper argues that S = X × V should not be interpreted as a classical phase space. Instead, the augmented variables form the configuration space of a constrained Hamiltonian system associated with the second-order control problem.

  • Control objective: The acceleration cost 1/2∥a∥2dt is not a classical physical action because classical actions typically exclude second-order time derivatives.The paper offers an alternative physical interpretation for this control objective.
  • Hamiltonian representation: The optimal second-order control system is equivalent to a Hamiltonian system in 4d-dimensional phase space, with q1 and q2 corresponding to x and v.The momentum p2 corresponds to the acceleration a.
  • Interpretation of S: Although S = X × V concatenates position and velocity, it is the configuration space of the classical mechanical system rather than its phase space.The associated Hamiltonian uses canonical coordinates and momenta in a larger phase space.
  • Constraint: The Lagrangian cannot be expressed solely through canonical coordinates and their first derivatives because the p1 term cannot be eliminated.The paper identifies this as characteristic of a constrained system.
  • Metric structure: The minimum transition cost on S is not a metric, and the configuration space cannot be equipped with the corresponding metric form.This follows because the Lagrangian lacks the required classical form.
Loading 2608.21070v1…