Trajectory Inference via Mean-field Langevin in Path Space

Lénaïc Chizat, Stephen Zhang, Matthieu Heitz, Geoffrey Schiebinger

Keywords: trajectory inference, path space, mean-field dynamics, interacting particle methods, Langevin algorithm, Wiener measure, entropic regularization, optimal transport, Schrödinger bridge

TL;DR: The estimator for trajectory inference that minimizes the entropy relative to Wiener measure can be computed with a Langevin dynamics in path space (convergence guaranteed).

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 in [Lavenant et al., 2021], 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.

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2205.07146/code)