Wasserstein Residuals: Learning Gradient Flows from Population Dynamics

Reconstructing the evolution of a population from sparse, discrete snapshots is a fundamental challenge in fields ranging from single-cell biology to crowd dynamics. These dynamics are often modeled as Wasserstein gradient flows (WGFs), which describe the steepest descent of an energy functional. Existing approaches, primarily based on the Jordan-Kinderlehrer-Otto (JKO) scheme, often struggle with long temporal gaps between observations and require solving costly optimal transport problems at each step. This paper addresses the inverse problem of recovering the underlying energy functional and the continuous trajectory of the population from observed snapshots.

Instead of relying on the JKO proximal scheme, the authors propose a residual-based framework. They define a non-negative loss function that vanishes exactly when the population's velocity field matches the gradient of a learned energy functional. This framework unifies existing paradigms like Path-Finding and Action Matching. The authors introduce stitching, a method that treats the population trajectory as a learnable particle cloud. By parameterizing the trajectory using a kernel density estimate (KDE) and coupling it with a data-fitting divergence, the method allows for global optimization of the entire trajectory without the need for simulation or costly optimal transport solvers.

The stitching method demonstrates state-of-the-art performance across several benchmarks, including single-cell RNA trajectory inference and the recovery of complex interaction dynamics. Unlike JKO-based methods, which often fail or collapse when snapshots are decorrelated or temporal gaps are large, stitching remains robust. The authors show that their method successfully recovers both potential and interaction kernels in synthetic experiments, capturing phase transitions in particle systems that other methods miss. Furthermore, the framework is flexible enough to extend beyond gradient flows to more general dynamical systems, such as chiral active matter.