[Paper Review] Trajectory Inference via Mean-field Langevin in Path Space

Lénaïc Chizat, Stephen X. Zhang+2 more|arXiv (Cornell University)|May 14, 2022

Single-cell and spatial transcriptomics4 citations

TL;DR

This paper proposes a grid-free, mean-field Langevin (MFL) dynamics algorithm for trajectory inference that leverages Schrödinger bridges and noisy gradient descent to compute the min-entropy estimator relative to the Wiener measure in path space. The method achieves global convergence to the true estimator with end-to-end theoretical guarantees, enabling consistent inference from sparse, time-resolved data with applications to single-cell RNA-seq and mass-variation modeling.

ABSTRACT

Trajectory inference aims at recovering the dynamics of a population from snapshots of its temporal marginals. To solve this task, a min-entropy estimator relative to the Wiener measure in path space was introduced by Lavenant et al. arXiv:2102.09204, and shown to consistently recover the dynamics of a large class of drift-diffusion processes from the solution of an infinite dimensional convex optimization problem. In this paper, we introduce a grid-free algorithm to compute this estimator. Our method consists in a family of point clouds (one per snapshot) coupled via Schrödinger bridges which evolve with noisy gradient descent. We study the mean-field limit of the dynamics and prove its global convergence to the desired estimator. Overall, this leads to an inference method with end-to-end theoretical guarantees that solves an interpretable model for trajectory inference. We also present how to adapt the method to deal with mass variations, a useful extension when dealing with single cell RNA-sequencing data where cells can branch and die.

Motivation & Objective

• To develop a grid-free, scalable algorithm for computing the min-entropy estimator in path space, which is otherwise computationally intractable due to infinite-dimensional convex optimization.
• To provide end-to-end theoretical guarantees for trajectory inference by proving global convergence of the proposed MFL dynamics to the true min-entropy estimator.
• To extend the method to handle mass variations in the data, such as cell branching and death, which are common in single-cell RNA-sequencing applications.
• To offer a theoretically grounded, interpretable alternative to non-convex deep learning-based trajectory inference methods that lack rigorous convergence guarantees.

Proposed method

• The method employs a family of point clouds (one per time snapshot) that are coupled via entropy-regularized optimal transport, or Schrödinger bridges, to model the dynamics over time.
• It formulates the inference problem as a mean-field Langevin dynamics in path space, where particles evolve via noisy gradient descent on a combined objective of data fitting and entropy regularization.
• The dynamics are governed by a stochastic differential equation that balances fidelity to observed marginals and smoothness via relative entropy to the Wiener measure.
• The algorithm uses annealing schedules for the step size $\eta$, noise level $\sigma$, and temperature $\tau$ to ensure stable convergence.
• A modified version of the dynamics is introduced to handle mass variations by incorporating a known growth rate $g(t,x)$, enabling modeling of cell proliferation and death.
• The method is implemented in a CPU-based framework with no explicit boundary enforcement, relying on particle stability within the support of observed data.

Experimental results

Research questions

• RQ1 Can a grid-free, stochastic optimization method be designed to compute the min-entropy estimator in path space with theoretical convergence guarantees?
• RQ2 How does the mean-field Langevin dynamics converge globally to the true trajectory inference solution, and what are the rates of convergence?
• RQ3 Can the method be extended to handle mass variations such as cell branching and death in single-cell data, which are not captured by standard transport models?
• RQ4 How does the performance of the proposed MFL method compare to existing methods like Global Waddington-OT in terms of reconstruction accuracy and robustness to sparse sampling?

Key findings

• The MFL dynamics converge globally to the min-entropy estimator at an exponential rate, establishing strong theoretical guarantees for trajectory inference.
• The method achieves state-of-the-art performance on synthetic and real single-cell data, with Energy Distance to the ground truth being minimized at $\lambda = 0.025$ for MFL and $\lambda = 0.01$ for gWOT.
• The algorithm successfully reconstructs complex dynamics such as bifurcating processes and time-dependent potentials, even with only 50 particles per timepoint.
• The inclusion of a known growth rate $g(t,x)$ enables accurate modeling of mass variations, such as cell proliferation, improving reconstruction fidelity in biologically realistic scenarios.
• The method remains stable and does not require explicit boundary conditions, as particles remain within the support of the observed data throughout optimization.
• The approach is robust to hyperparameter choices and achieves consistent performance across multiple datasets and parameter settings, as validated through energy distance and visual reconstruction quality.