Generalized Moment-Informed Flow Matching for Trajectory Inference

Trajectory inference seeks to reconstruct continuous dynamics from sparse, unpaired snapshots of evolving populations. These observations do not uniquely determine the underlying path process, so modeling assumptions are needed to select plausible dynamics. Collecting more samples per snapshot or observing additional time points can improve reconstruction, but detailed measurements are often costly or destructive. We therefore consider settings in which less costly measurements provide aggregate information in the form of expectations of test functions under intermediate marginal distributions. We build on flow matching, which learns a vector field to reproduce the marginal evolution of a prescribed stochastic interpolant through simulation-free regression. To incorporate the aggregate observations into this framework, we introduce Generalized Moment-Informed Flow Matching (GMI-FM). Our method first learns an endpoint-preserving neural interpolant under the aggregate constraints using an augmented-Lagrangian formulation, then regresses the corresponding vector field against the interpolant’s velocities, keeping both stages simulation-free. Our framework also accommodates geometric priors, such as data-dependent metrics, through the interpolant regularization, encouraging learned trajectories to follow the geometry of the data. We provide theoretical grounding for GMI-FM by relating its construction to a Schrödinger bridge formulation with the same aggregate constraints. Experiments on synthetic and real-world datasets demonstrate that GMI-FM achieves better performance on marginal reconstruction metrics than competing methods. These findings highlight the potential of integrating lower-cost aggregate measurements with costly, detailed population snapshots to improve trajectory inference.