Source-linked AI summary
TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics
Alexander Tong, Jessie Huang, Guy Wolf, David van Dijk, Smita Krishnaswamy
TL;DR
Static cross-sectional measurements make it difficult to infer individual cellular dynamics because entities are not observed longitudinally, while existing approaches do not capture continuous nonlinear paths. TrajectoryNet links continuous normalizing flows to dynamic optimal transport and uses Neural ODEs to model smooth trajectories across distributions. In cellular and artificial experiments, it improves interpolation of intermediate distributions over baseline approaches, while its formulation includes assumptions and regularization that do not guarantee mass conservation and currently uses direction rather than magnitude for velocity information.
Problem
Cross-sectional timepoint data lacks individual correspondences and existing interpolation methods do not model continuous nonlinear paths between distributions.
Method
TrajectoryNet links continuous normalizing flows with dynamic optimal transport and uses a Neural ODE to learn smooth continuous-time paths across distributions.
Results
TrajectoryNet improves interpolation of held-out cellular distributions over baseline models and captures curved-manifold structure that straight-path methods miss.
Takeaways & Limitations
TrajectoryNet provides continuous-time trajectories for cellular systems whose measurements are static snapshots and can incorporate information from multiple timepoints.
Takeaways & Limitations
Growth-rate regularization does not guarantee conservation of mass, and current velocity estimates provide direction but not reliable magnitude.
Abstract
from arXiv · showhide
It is increasingly common to encounter data from dynamic processes captured by static cross-sectional measurements over time, particularly in biomedical settings. Recent attempts to model individual trajectories from this data use optimal transport to create pairwise matchings between time points. However, these methods cannot model continuous dynamics and non-linear paths that entities can take in these systems. To address this issue, we establish a link between continuous normalizing flows and dynamic optimal transport, that allows us to model the expected paths of points over time. Continuous normalizing flows are generally under constrained, as they are allowed to take an arbitrary path from the source to the target distribution. We present TrajectoryNet, which controls the continuous paths taken between distributions to produce dynamic optimal transport. We show how this is particularly applicable for studying cellular dynamics in data from single-cell RNA sequencing (scRNA-seq) technologies, and that TrajectoryNet improves upon recently proposed static optimal transport-based models that can be used for interpolating cellular distributions.
1. Introduction
TrajectoryNet addresses the challenge of inferring continuous, nonlinear dynamics from cross-sectional distributions without point-to-point correspondences. It links continuous normalizing flows with dynamic optimal transport and applies the resulting model to cellular trajectories.
- Motivation: Cross-sectional measurements provide distributions at multiple timepoints but usually lack longitudinal correspondences between individual entities.This makes interpolation and recovery of development or disease dynamics difficult.
- Method: TrajectoryNet formulates the problem as unbalanced dynamic transport, moving entities between cross-sectional measurements along efficient, smooth paths.Its formulation connects continuous normalizing flows with dynamic optimal transport through a Neural ODE framework.
- Method: Continuous normalizing flows extend optimal transport to multiple, possibly unevenly spaced distributions and smooth flows in high dimensions.This setting matches time-series data from single-cell RNA sequencing.
- Application: Single-cell RNA sequencing reveals detailed cellular states but destroys cells during measurement, preventing direct observation of individual changes over time.Only a handful of discrete timepoints are generally collected because the technology is expensive.
- Advantages: TrajectoryNet interpolates along the manifold of observed entities and creates continuous-time trajectories for individual entities.It is presented as a nonlinear alternative to methods that infer within-timepoint trajectories or interpolate linearly between two timepoints.
- Scope: The proposed penalties are also applicable to modeling dynamics from cross-sectional population data beyond cellular systems.The paper specifically notes potential use in other settings where population-level snapshots are used to model dynamics.
2. Background and Related Work
Prior work includes static and dynamic optimal-transport methods, single-cell pseudotime approaches, and distributional splines, but these methods either simplify between-timepoint motion or lack efficient high-dimensional implementations. TrajectoryNet combines global continuous-time inference across multiple timepoints with local dynamics.
- Optimal Transport: Static optimal transport can be computationally expensive, with the discrete linear-program formulation costing O(n3), motivating faster approximations.Entropic regularization and Sinkhorn-based methods provide parallelizable approximations in discrete settings.
- Dynamic Optimal Transport: Dynamic optimal transport introduces time interpolation and fluid-dynamics links, but grid-based solvers scale exponentially with dimension.TrajectoryNet contributes an approximate neural-network solver for high-dimensional smooth problems.
- Single-Cell Trajectories: Single-cell pseudotime methods infer developmental axes within one snapshot and depend on prior assumptions about data structure.They do not directly provide trajectories across multiple measured timepoints.
- Multiple Timepoints: Methods using multiple timepoints generally match coarse-grained populations without describing motion between measured timepoints, often assuming minimum-Euclidean-cost paths.Such straight-path assumptions may be unrealistic for the underlying dynamics.
- TrajectoryNet: TrajectoryNet unifies these directions by inferring global continuous-time trajectories from multiple timepoints while respecting local dynamics within each timepoint.Its Figure 1 uses a Neural ODE to learn dynamics derivatives and integrates between sampled distributions.
3. Preliminaries
The paper reviews static and dynamic optimal transport alongside continuous normalizing flows, then connects CNFs to high-dimensional dynamic optimal transport. This connection provides a computational route for modeling distributions and trajectories over time.
- The Monge-Kantorovich Problem: Static optimal transport compares source and target distributions but can be computationally expensive for discrete and high-dimensional measures.Discrete static OT has cost O(n^3), while continuous high-dimensional settings remain difficult to scale.
- Dynamic Optimal Transport: Dynamic optimal transport introduces time-dependent density and velocity fields constrained by the continuity equation.The velocity field transports mass while preventing points in the field from acting as sources or sinks.
- Dynamic Optimal Transport: The dynamic OT solution minimizes the L2 norm of a velocity field that transports the source distribution to the target distribution.This formulation yields a minimum-cost flow and linearly interpolates between the endpoint measures.
- Continuous Normalizing Flows: Continuous normalizing flows transform an initial distribution continuously by integrating a time-dependent differential equation.Unlike a static mapping, a CNF models the measure over time through the velocity field fθ(x, t).
- Continuous Normalizing Flows: TrajectoryNet exploits the connection between CNFs and dynamic OT to provide a high-dimensional solution to the dynamic optimal transport problem.The approach is motivated by the poor scaling of grid-based dynamic OT methods in high dimensions.
4. TrajectoryNet: Efficient Dynamic Optimal Transport
TrajectoryNet adapts continuous normalizing flows to approximate dynamic optimal transport, adding energy and biological regularizations for smooth, cellularly plausible paths. The framework supports unbalanced transport, manifold-aware trajectories, and known local velocities.
- Dynamic OT Approximation via Regularized CNF: TrajectoryNet adds energy regularization to a continuous normalizing flow so its paths approximate dynamic optimal transport.The regularization penalizes the norm of the flow and relaxes endpoint matching through a divergence penalty.
- Dynamic OT Approximation via Regularized CNF: Energy regularization produces straighter paths similar to the optimal map, but excessive regularization can undershoot the target distribution.A Jacobian penalty or additional training noise helps obtain straight paths with lower energy regularization.
- Dynamic OT Approximation via Regularized CNF: A single smooth function models the full time series, avoiding discontinuities that arise when optimal transport is performed only between successive sampled timepoints.The resulting flow is described as the minimal-cost smooth flow over time.
- Biological Priors: TrajectoryNet incorporates growth-rate regularization to model unbalanced transport in systems characterized by growth and death.The implementation adapts discrete static unbalanced optimal transport to the continuous setting using a learned growth-rate network.
- Biological Priors: Density and velocity regularizations encourage paths to follow dense manifold regions and match local estimates of cellular velocity.The density penalty uses nearby measured points, while the velocity penalty aligns the flow derivative with observed dx/dt estimates.
- Biological Priors: Growth-rate regularization does not guarantee conservation of mass.Normalizing the learned mass distribution during training would require an integration over the ambient space that is too computationally costly.
5. Experiments
Experiments evaluate TrajectoryNet on artificial manifolds and single-cell datasets by interpolating held-out timepoints and comparing predicted distributions or trajectories with baselines. The results show that regularization helps recover nonlinear manifolds and that TrajectoryNet performs well on cellular dynamics.
- Evaluation: Held-out timepoints are evaluated using EMD, for which lower values indicate more accurate interpolations.Artificial trajectories are additionally evaluated with MSE against known paths.
- Artificial Data: On curved 1D manifolds, density or velocity regularization makes TrajectoryNet follow the data manifold instead of the shortest straight path.The artificial datasets include arch and tree structures with overlapping sampled timepoints.
- Artificial Data: The circle dataset demonstrates that velocity estimates are necessary to recover an unchanging distribution whose points move counterclockwise.Without velocity estimates, this dynamical system cannot be identified.
- Single-Cell Data: TrajectoryNet outperforms baseline models on the mouse cortex dataset, especially when density and velocity information are added.The dataset spans mouse embryo days E12.5, E14.5, E16, and E17.5, with neuron development as its major variation axis.
- Single-Cell Data: The Embryoid body dataset contains five timepoints in which stem cells differentiate into roughly four precursor types with a branching structure.Velocity regularization does not appear to help, possibly because the data have low unspliced RNA counts.
- Single-Cell Data: TrajectoryNet trajectories can be projected back into gene space, where paths sampled from four terminal branches recapitulate known biology.The trajectories are integrated backwards from curated endpoint cells.
6. Conclusion
TrajectoryNet computes dynamic optimal transport between discrete-time sample distributions to model realistic continuous paths. The paper argues that this is useful for high-dimensional, dynamic, nonlinear systems, while existing alternatives are either under-powered or computationally infeasible.
- Conclusion: TrajectoryNet models realistic continuous-time sample paths between distributions observed at discrete times.In single-cell data, it reconstructs trajectories for cells destroyed during measurement.
- Conclusion: The framework targets high-dimensional, dynamic, nonlinear systems where static optimal transport does not interpolate well between measured timepoints.The conclusion specifically describes this application to single-cell measurements.
- Conclusion: Existing non-neural dynamic optimal transport methods are described as computationally infeasible for this task.
- Conclusion: The authors integrate multiple priors and assumptions to bias TrajectoryNet toward more realistic dynamic optimal transport solutions.They identify stochastic dynamics and jointly learning growth with dynamics as future directions.
A.1. Proof of Theorem 4.1
The proof analyzes a Lagrangian relaxation of the constrained optimization problem. It establishes monotonic behavior of the objective and divergence terms and shows convergence to the constrained optimum.
- Lagrangian Relaxation: The proof introduces a Lagrange multiplier λ for the minimization problem subject to the stated constraints.The initial distribution satisfies zero KL divergence because the flow starts from the base distribution.
- Monotonicity: The divergence term decreases monotonically as λ increases and converges to zero as λ approaches infinity.
- Monotonicity: The relaxed objective increases monotonically with λ because the objective term and divergence penalty have complementary monotonic behavior.
- Convergence: The proof constructs an accelerated transformation followed by a correction to achieve zero divergence while increasing the energy term.This construction supports the limiting argument for the relaxed problem.
- Convergence: For sufficiently large λ, the max-min solution converges to the solution of the original constrained problem.The argument uses upper and lower bounds on the relaxed objective.
B. Growth Rate Model Training
The growth model trains a neural network to reproduce growth rates from a discrete unbalanced optimal transport solution while extending them continuously across measured timepoints.
- Model Training: The growth network is trained to match a discrete unbalanced optimal transport problem with entropic regularization.
- Model Training: The discrete solution provides a growth rate for each source cell in the transition from μ to ν.
- Model Training: A neural network learns continuous growth rates by matching measured-cell rates and assigning one to uniformly sampled negative points.The implementation uses equally sized batches of real and uniformly sampled points.
C. Scaling with Dimension
TrajectoryNet is designed to scale polynomially with data dimensionality, unlike grid-based dynamic optimal transport methods whose grid size grows exponentially. Experiments report roughly linear seconds per function evaluation as dimensionality increases, without implying linear convergence.
- Scaling with Dimension: Grid-based dynamic optimal transport methods scale exponentially with dimension because their discretized grid gains exponentially many points.These methods are therefore applicable mainly in low dimensions.
- Scaling with Dimension: TrajectoryNet scales polynomially with dimension.
- Scaling with Dimension: Seconds per evaluation are roughly linear in data dimensionality on the mouse cortex computation-time experiment.The experiment used 100 batches of 1000 points across different dimensionalities.
D. Biological Considerations
The biological analyses use regularization and trajectory projections to study cellular systems with manifold structure, local velocity information, and growth or differentiation. RNA-velocity quality is an important constraint, particularly for the Embryoid body dataset's low unspliced-count fraction.
- Biological Considerations: Low unspliced RNA counts in the Embryoid body dataset make RNA-velocity estimates noisier and lower quality.Each timepoint contains around 10%-20% unspliced RNA, relatively low compared with other recent works.
- Biological Considerations: Density and velocity regularization can be used to follow a 1D manifold in 2D.This behavior is illustrated for a curved manifold setting.
- Biological Considerations: The Embryoid body analysis compares spliced, ambiguous, and unspliced RNA-count ratios across five timepoints.
- Biological Considerations: TrajectoryNet projections into gene space can identify progenitor features earlier than they are identifiable in gene space alone.The analysis reports HAND1 as high at earlier timepoints and potentially able to distinguish the population earlier.
- Biological Considerations: The synthetic cycle dataset assigns points tangent velocity of magnitude π/5, so points move one-tenth around the circle per time unit.
E.2. Single Cell Datasets
The study analyzes two 10X single-cell datasets processed into spliced and unspliced count matrices without batch correction. Parameter selection used a limited grid search, and the reported software environment is specified explicitly.
- E.2. Single Cell Datasets: Both single-cell datasets were generated with 10X sequencing and processed with CellRanger and velocyto.The Embryoid body dataset contains roughly 30,000 unfiltered and 16,000 filtered cells; the mouse cortex dataset contains roughly 20,000 filtered cells.
- E.2. Single Cell Datasets: No batch correction was used for either single-cell dataset.
- E.2. Single Cell Datasets: The regularization parameters were selected by grid search over λdensity ∈ {0, 0.1, 0.01} and λvelocity ∈ {0, 0.001, 0.0001}.
- E.2. Single Cell Datasets: The authors state that a more extensive parameter search could improve results and that the search was intended to demonstrate approach viability rather than fully explore parameter space.
- E.2. Single Cell Datasets: The reported implementation uses scvelo 0.1.24, torch 1.3.1, torchdiffeq 0.0.1, velocyto 0.17.17, and scanpy 1.4.5.