Partially Observed Trajectory Inference using Optimal Transport and a Dynamics Prior

TL;DR

This paper introduces a novel method for inferring latent trajectories from partial observations using optimal transport and dynamics priors, extending existing models to latent SDEs with theoretical guarantees and practical robustness.

Contribution

It extends trajectory inference to latent SDEs with partial observations, introducing the PO-MFL algorithm and providing theoretical and empirical validation.

Findings

The PO-MFL algorithm effectively infers latent trajectories from partial data.

The method demonstrates exponential convergence and robustness in experiments.

It significantly outperforms baseline methods in key scenarios.

Abstract

Trajectory inference seeks to recover the temporal dynamics of a population from snapshots of its (uncoupled) temporal marginals, i.e. where observed particles are not tracked over time. Prior works addressed this challenging problem under a stochastic differential equation (SDE) model with a gradient-driven drift in the observed space, introducing a minimum entropy estimator relative to the Wiener measure and a practical grid-free mean-field Langevin (MFL) algorithm using Schr\"odinger bridges. Motivated by the success of observable state space models in the traditional paired trajectory inference problem (e.g. target tracking), we extend the above framework to a class of latent SDEs in the form of observable state space models. In this setting, we use partial observations to infer trajectories in the latent space under a specified dynamics model (e.g. the constant velocity/acceleration models from target tracking). We introduce the PO-MFL algorithm to solve this latent trajectory inference problem and provide theoretical guarantees to the partially observed setting. Experiments validate the robustness of our method and the exponential convergence of the MFL dynamics, and demonstrate significant outperformance over the latent-free baseline in key scenarios.