Longitudinal Flow Matching for Trajectory Modeling
Mohammad Mohaiminul Islam, Thijs P. Kuipers, Sharvaree Vadgama, Coen de Vente, Afsana Khan, Clara I. Sánchez, Erik J. Bekkers
TL;DR
IMMFM tackles the problem of modeling high-dimensional, sparsely observed longitudinal trajectories by learning a continuous stochastic flow that aligns with multiple timepoints. It introduces a piecewise-quadratic conditional path and jointly learns drift $v_\theta$ and diffusion $g_\theta$, augmented by a data-driven diffusion to capture uncertainty, with MMOT-based handling of incomplete trajectories. The method achieves superior forecasting and downstream task performance on synthetic benchmarks and real longitudinal neuroimaging data (e.g., ADNI, GBM, MS), with notable gains from the uncertainty-aware SU-IMMFM variant and 3D extensions. The work advances subject-specific trajectory modeling under irregular sampling, offering practical impact for prognosis and treatment planning in clinical settings, and opens directions toward temporally aware latent spaces and physics-informed constraints.
Abstract
Generative models for sequential data often struggle with sparsely sampled and high-dimensional trajectories, typically reducing the learning of dynamics to pairwise transitions. We propose Interpolative Multi-Marginal Flow Matching (IMMFM), a framework that learns continuous stochastic dynamics jointly consistent with multiple observed time points. IMMFM employs a piecewise-quadratic interpolation path as a smooth target for flow matching and jointly optimizes drift and a data-driven diffusion coefficient, supported by a theoretical condition for stable learning. This design captures intrinsic stochasticity, handles irregular sparse sampling, and yields subject-specific trajectories. Experiments on synthetic benchmarks and real-world longitudinal neuroimaging datasets show that IMMFM outperforms existing methods in both forecasting accuracy and further downstream tasks.