Table of Contents
Fetching ...

Partially Observed Trajectory Inference using Optimal Transport and a Dynamics Prior

Anming Gu, Edward Chien, Kristjan Greenewald

TL;DR

The paper tackles partially observed trajectory inference by coupling latent state dynamics with unseen velocities via a latent SDE and observable map. It introduces PO-MFL, a minimum-entropy estimator that links time-point marginals through entropic OT under a dynamics prior, and solves it with a mean-field Langevin algorithm, providing exponential-convergence guarantees. The reduced formulation over latent marginals and a tractable Euler–Maruyama OT cost enable sampling of latent trajectory distributions and, hence, trajectory samples in the observation space. Empirical results on synthetic and real data show that incorporating a dynamics prior substantially improves robustness and accuracy over latent-free baselines, with theoretical consistency guarantees supporting the approach.

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ödinger 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.

Partially Observed Trajectory Inference using Optimal Transport and a Dynamics Prior

TL;DR

The paper tackles partially observed trajectory inference by coupling latent state dynamics with unseen velocities via a latent SDE and observable map. It introduces PO-MFL, a minimum-entropy estimator that links time-point marginals through entropic OT under a dynamics prior, and solves it with a mean-field Langevin algorithm, providing exponential-convergence guarantees. The reduced formulation over latent marginals and a tractable Euler–Maruyama OT cost enable sampling of latent trajectory distributions and, hence, trajectory samples in the observation space. Empirical results on synthetic and real data show that incorporating a dynamics prior substantially improves robustness and accuracy over latent-free baselines, with theoretical consistency guarantees supporting the approach.

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ödinger 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.
Paper Structure (35 sections, 27 theorems, 114 equations, 10 figures, 1 table, 2 algorithms)

This paper contains 35 sections, 27 theorems, 114 equations, 10 figures, 1 table, 2 algorithms.

Key Result

Theorem 3.1

Suppose $\mathbf{P}$ is the SDE eq:SDE with initial condition $\mathbf{P}_0\in\mathcal{P}(\mathcal{X})$ such that $H(\mathbf{P}_0|\mathrm{vol})<+\infty$. Let $\mathbf{R}^{T,\lambda,\sigma}\in\mathcal{P}(\Omega)$ be the unique minimizer of eq:functional_path_space: If $\{t_i^T\}_{t\in [T]}$ becomes dense in $[0, 1]$, then $\lim_{\sigma\to 0,\lambda \to 0} \left(\lim_{T \to \infty} \mathbf{R}^{T,\l

Figures (10)

  • Figure 1: Constant velocity model, where the variance of the ground truth has been rescaled. We see that our method, PO-MFL, is more robust as the baseline method, FO-MFL, fails to converge, and provides per-particle velocity in contrast to FO-MFL. See Section \ref{['sec:experiments']} for the experiment setting.
  • Figure 2: (left) Velocity of one particle at end of optimization. (right) Population velocity at beginning of optimization, showing exponential convergence.
  • Figure 3: Average $W_2$ distance between ground truth and PO-MFL (blue) and FO-MFL (orange) recovered positions in "constant velocity" model (App. \ref{['sec:additional_experiments']}). (L) Number of time points. (C) Number of observations. (R) Velocity.
  • Figure 4: Crossing paths experiment under the "constant velocity" SDE.
  • Figure 5: Wikipedia page traffic data.
  • ...and 5 more figures

Theorems & Definitions (51)

  • Definition 1: $\mathcal{C}_\Psi$-ensemble observability
  • Theorem 3.1: Consistency (informal, see Thm. \ref{['thm:2.3']})
  • Theorem 3.2: Representer theorem
  • Proposition 3.3
  • Theorem 3.4: Convergence
  • Definition 2: Ensemble observability zeng2015ensemble_observability
  • Proposition B.2: zeng2015ensemble_observability
  • Corollary B.3: zeng2015ensemble_observability
  • Corollary B.4
  • proof
  • ...and 41 more