Trajectory Inference via Mean-field Langevin in Path Space
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\"odinger 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.